fix(parser): preserve map iterator schemas

Author: fffonion

src/compiler/parser/statements.rs

@@ -1241,11 +1241,21 @@ impl Parser {
 
         let previous_key_slot = self.replace_current_local_binding(&key_name, key_slot);
         let previous_value_slot = self.replace_current_local_binding(&value_name, value_slot);
+        let previous_key_schema = self.local_schemas.get(&key_slot).cloned();
+        let previous_value_schema = self.local_schemas.get(&value_slot).cloned();
+        if let Some(schema) = key_schema.as_ref() {
+            self.local_schemas.insert(key_slot, schema.clone());
+        }
+        if let Some(schema) = value_schema.as_ref() {
+            self.local_schemas.insert(value_slot, schema.clone());
+        }
         self.borrowed_map_iter_locals.push(map_slot);
         self.loop_depth += 1;
         let body_result = self.parse_block("expected '{' after borrowed map");
         self.loop_depth -= 1;
         self.borrowed_map_iter_locals.pop();
+        self.restore_local_schema(value_slot, previous_value_schema);
+        self.restore_local_schema(key_slot, previous_key_schema);
         self.restore_current_local_binding(&value_name, previous_value_slot);
         self.restore_current_local_binding(&key_name, previous_key_slot);
         let mut body = body_result?;
@@ -1328,6 +1338,14 @@ impl Parser {
         }
     }
 
+    fn restore_local_schema(&mut self, slot: LocalSlot, previous: Option<TypeSchema>) {
+        if let Some(schema) = previous {
+            self.local_schemas.insert(slot, schema);
+        } else {
+            self.local_schemas.remove(&slot);
+        }
+    }
+
     fn bind_for_loop_local(&mut self, name: &str) -> Result<(LocalSlot, bool), ParseError> {
         if !self.closure_scopes.is_empty() {
             if let Some(index) = self

tests/example_tests.rs

@@ -387,12 +387,10 @@ fn rustscript_untyped_rebinding_does_not_reuse_stale_map_schema() {
     let source = r#"
 let values: map<int> = {a: 1};
 let values = {a: "x"};
-for (key: string, value: int) in &values {
-    value;
-}
+for (key: string, value: int) in &values { value; }
 "#;
     let err = match compile_source_with_flavor(source, SourceFlavor::RustScript) {
-        Ok(_) => panic!("untyped rebinding must not retain the old map<int> schema"),
+        Ok(_) => panic!("untyped rebinding must clear stale map schema"),
         Err(err) => err,
     };
     assert!(err.to_string().contains("map<T> schema"));
@@ -403,9 +401,7 @@ fn rustscript_function_parameter_map_type_is_visible_to_iterators() {
     let source = r#"
 fn sum(values: map<int>) -> int {
     let mut total: int = 0;
-    for (key: string, value: int) in &values {
-        total += value;
-    }
+    for (key: string, value: int) in &values { total += value; }
     total
 }
 sum({a: 1, b: 2});
@@ -422,9 +418,7 @@ fn rustscript_borrowed_map_iterator_propagates_nested_binding_schema() {
 let groups: map<map<int>> = {first: {a: 1, b: 2}};
 let mut total: int = 0;
 for (group_key: string, group: map<int>) in &groups {
-    for (item_key: string, item: int) in &group {
-        total += item;
-    }
+    for (item_key: string, item: int) in &group { total += item; }
 }
 total;
 "#;