From 0e52ed7a49252645acaac01cb4171d88db9f8788 Mon Sep 17 00:00:00 2001
From: Jeff <dev@jeffa.io>
Date: Tue, 2 Jul 2024 21:06:23 -0400
Subject: [PATCH] Improve context recursion

---
 dust-lang/src/abstract_tree/mod.rs |  1 -
 dust-lang/src/context.rs           | 11 ++++++++---
 2 files changed, 8 insertions(+), 4 deletions(-)

diff --git a/dust-lang/src/abstract_tree/mod.rs b/dust-lang/src/abstract_tree/mod.rs
index a90460c..ac58748 100644
--- a/dust-lang/src/abstract_tree/mod.rs
+++ b/dust-lang/src/abstract_tree/mod.rs
@@ -58,7 +58,6 @@ pub use self::{
 use crate::{
     context::Context,
     error::{DustError, RuntimeError, ValidationError},
-    identifier::Identifier,
     Value,
 };
 
diff --git a/dust-lang/src/context.rs b/dust-lang/src/context.rs
index 791414d..f96c951 100644
--- a/dust-lang/src/context.rs
+++ b/dust-lang/src/context.rs
@@ -219,10 +219,15 @@ impl Context {
     pub fn set_type(&self, identifier: Identifier, r#type: Type) -> Result<(), PoisonError> {
         log::debug!("Setting {identifier} to type {}", r#type);
 
-        self.data
-            .write()?
+        let mut data = self.data.write()?;
+        let usage_data = data
             .variables
-            .insert(identifier, (VariableData::Type(r#type), UsageData::new()));
+            .remove(&identifier)
+            .map(|(_, usage_data)| usage_data)
+            .unwrap_or_else(|| UsageData::new());
+
+        data.variables
+            .insert(identifier, (VariableData::Type(r#type), usage_data));
 
         Ok(())
     }