Replace finding each module in the source code with iterating over co…

…de directly

src/compiler_top.rs

@@ -5,7 +5,7 @@ use tree_sitter::Parser;
 use crate::{
     errors::ErrorCollector,
     file_position::FileText,
-    flattening::{flatten, initialization::gather_initial_file_data, typechecking::typecheck_all_modules},
+    flattening::{flatten_all_modules, initialization::gather_initial_file_data, typechecking::typecheck_all_modules},
     instantiation::InstantiatedModule,
     linker::{FileData, FileUUID, Linker, ModuleUUID}
 };
@@ -45,20 +45,11 @@ pub fn update_file(text : String, file_id : FileUUID, linker : &mut Linker) {
 }
 
 pub fn recompile_all(linker : &mut Linker) {
-    // Flatten all modules
-    let id_vec : Vec<ModuleUUID> = linker.modules.iter().map(|(id, _)| id).collect();
-    for id in id_vec {
-        //let md = &linker.modules[id];// Have to get them like this, so we don't have a mutable borrow on self.modules across the loop
-
-        flatten(linker, id);
-
-        let md = &mut linker.modules[id]; // Convert to mutable ptr
-        md.instantiations.clear_instances();
-    }
-
+    flatten_all_modules(linker);
     typecheck_all_modules(linker);
 
-    // Can't merge these loops, because instantiation can only be done once all modules have been type checked
+    // Make an initial instantiation of all modules
+    // Won't be possible once we have template modules
     for (id, _md) in &linker.modules {
         //md.print_flattened_module();
         // Already instantiate any modules without parameters

src/flattening/mod.rs

@@ -12,7 +12,7 @@ use crate::{
     errors::{error_info, ErrorCollector, ErrorInfo},
     file_position::{BracketSpan, Span},
     instantiation::InstantiationList,
-    linker::{ConstantUUID, FileUUID, LinkInfo, Linker, ModuleEditContext, ModuleUUID, NameElem, NamedConstant, NamedType, ResolvedName, TypeUUIDMarker},
+    linker::{ConstantUUID, FileData, FileUUID, LinkInfo, Linker, ModuleEditContext, ModuleUUID, NameElem, NamedConstant, NamedType, ResolvedName, TypeUUIDMarker},
     parser::{Cursor, Documentation},
     typing::{get_binary_operator_types, typecheck, typecheck_is_array_indexer, typecheck_unary_operator, AbstractType, WrittenType, BOOL_TYPE, INT_TYPE},
     value::Value
@@ -453,28 +453,28 @@ impl<'l> LocalOrGlobal<'l> {
 
 
 
-struct FlatteningContext<'l> {
-    upper_ctx : ModuleEditContext<'l>,
+struct FlatteningContext<'md, 'l> {
+    upper_ctx : ModuleEditContext<'md, 'l>,
     ports_to_visit : UUIDRange<PortIDMarker>,
 
     local_variable_context : LocalVariableContext<'l, FlatID>
 }
 
-impl<'l> Deref for FlatteningContext<'l> {
-    type Target = ModuleEditContext<'l>;
+impl<'md, 'l> Deref for FlatteningContext<'md, 'l> {
+    type Target = ModuleEditContext<'md, 'l>;
 
     fn deref(&self) -> &Self::Target {
         &self.upper_ctx
     }
 }
 
-impl<'l> DerefMut for FlatteningContext<'l> {
+impl<'md, 'l> DerefMut for FlatteningContext<'md, 'l> {
     fn deref_mut(&mut self) -> &mut Self::Target {
         &mut self.upper_ctx
     }
 }
 
-impl<'l> FlatteningContext<'l> {
+impl<'md, 'l> FlatteningContext<'md, 'l> {
     fn resolve_identifier(&self, identifier_span : Span) -> LocalOrGlobal {
         // Possibly local
         let name_text = &self.file.file_text[identifier_span];
@@ -1135,28 +1135,54 @@ pub struct FlattenedInterfacePort {
     pub span : Span
 }
 
-/*
-This method flattens all given code into a simple set of assignments, operators and submodules. 
-It already does basic type checking and assigns a type to every wire. 
-The Generating Structure of the code is not yet executed. 
-It is template-preserving
-*/
-pub fn flatten(linker : &mut Linker, module_uuid : ModuleUUID) {
-    let upper_ctx = ModuleEditContext::new(linker, module_uuid);
+
+/// This method flattens all given code into a simple set of assignments, operators and submodules. 
+/// It already does basic type checking and assigns a type to every wire. 
+/// The Generating Structure of the code is not yet executed. 
+/// It is template-preserving
+/// 
+/// Separate 'md lifetime for the module. 
+/// For some reason if it has the same lifetime as the linker ('l), 
+/// then the compiler thinks we could store cursor elements in the module, which would be bad? 
+/// Don't fully understand this, but separating the lifetimes makes it work. 
+fn flatten<'md, 'l>(linker : *const Linker, md : &'md mut Module, file_data : &'l FileData, cursor : &mut Cursor<'l>) {
+    let upper_ctx = ModuleEditContext::new(linker, file_data, md);
     println!("Flattening {}", upper_ctx.md.link_info.name);
 
-    let mut context = FlatteningContext{
+    let mut context = FlatteningContext {
         ports_to_visit : upper_ctx.md.module_ports.ports.id_range(),
         upper_ctx,
         local_variable_context : LocalVariableContext::new_initial()
     };
 
-    // The given span should correspond perfectly to this, so impossible we don't find the node. 
-    let mut cursor = Cursor::new_for_node(&context.file.tree, &context.file.file_text, context.md.link_info.span, kind!("module"));
-
-    // Temporary, switch to iterating over nodes in file itself when needed. 
-    context.flatten_module(&mut cursor);
+    context.flatten_module(cursor);
 
     // Make sure all ports have been visited
     assert!(context.ports_to_visit.is_empty());
 }
+
+/// Flattens all modules in the project. 
+/// 
+/// Requires that first, all modules have been initialized. 
+pub fn flatten_all_modules<'l>(linker : &'l mut Linker) {
+    let linker_ptr : *const Linker = linker;
+    let modules : &'l mut ArenaAllocator<_,_> = &mut linker.modules;
+
+    for (_file_id, file) in &linker.files {
+        let mut associated_value_iter = file.associated_values.iter();
+
+        let mut cursor = Cursor::new_at_root(&file.tree, &file.file_text);
+
+        cursor.list(kind!("source_file"), |cursor| {
+            match cursor.kind() {
+                kind!("module") => {
+                    let Some(NameElem::Module(module_uuid)) = associated_value_iter.next() else {unreachable!()};
+
+                    let md : &mut Module = &mut modules[*module_uuid];
+                    flatten(linker_ptr, md, file, cursor);
+                }
+                other => todo!("{}", tree_sitter_sus::language().node_kind_for_id(other).unwrap())
+            }
+        });
+    }
+}

src/linker/mod.rs

@@ -160,6 +160,7 @@ impl Linkable for NamedType {
 pub struct FileData {
     pub file_text : FileText,
     pub parsing_errors : ErrorCollector,
+    /// In source file order
     pub associated_values : Vec<NameElem>,
     pub tree : tree_sitter::Tree
 }

src/linker/resolver.rs

@@ -143,26 +143,23 @@ impl<'l> ResolvedName<'l> {
 
 
 
-pub struct ModuleEditContext<'linker> {
+pub struct ModuleEditContext<'md, 'linker> {
     /// The module we are currently editing
-    pub md : &'linker mut Module,
+    pub md : &'md mut Module,
 
     pub file : &'linker FileData,
 
     resolver : UnsafeGlobalResolver
 }
 
-impl<'linker> ModuleEditContext<'linker> {
+impl<'md, 'linker> ModuleEditContext<'md ,'linker> {
     /// See [ModuleEditContext::drop]
-    pub fn new(linker : &'linker mut Linker, module_uuid : ModuleUUID) -> Self {
-        let linker_ptr : *const Linker = linker;
-        let md = &mut linker.modules[module_uuid];
-
+    pub fn new(linker_ptr : *const Linker, file : &'linker FileData, md : &'md mut Module) -> Self {
         Self {
-            file : &linker.files[md.link_info.file],
+            file,
             resolver : UnsafeGlobalResolver {
                 linker : linker_ptr,
-                file_text : &linker.files[md.link_info.file].file_text,
+                file_text : &file.file_text,
                 errors : md.link_info.errors.take(),
                 resolved_globals : RefCell::new(ResolvedGlobals::empty())
             },
@@ -205,15 +202,15 @@ impl<'linker> ModuleEditContext<'linker> {
 }
 
 /// Don't actually need [::core::ops::DerefMut]
-impl<'linker> Deref for ModuleEditContext<'linker> {
+impl<'md, 'linker> Deref for ModuleEditContext<'md, 'linker> {
     type Target = UnsafeGlobalResolver;
 
     fn deref(&self) -> &Self::Target {
         &self.resolver
     }
 }
 
-impl<'linker> Drop for ModuleEditContext<'linker> {
+impl<'md, 'linker> Drop for ModuleEditContext<'md, 'linker> {
     /// Places errors and resolved globals back in Module's LinkInfo
     /// 
     /// See [ModuleEditContext::new]

src/parser.rs

@@ -62,26 +62,6 @@ impl<'t> Cursor<'t> {
         Self{cursor : tree.walk(), file_text, gathered_comments : Vec::new(), current_field_was_already_consumed : false}
     }
 
-    pub fn new_for_node(tree : &'t Tree, file_text : &'t FileText, span : Span, kind : u16) -> Self {
-        let mut cursor = tree.walk();
-        assert!(cursor.goto_first_child());
-
-        let span_range = span.into_range();
-        loop {
-            let node = cursor.node();
-            if node.byte_range() == span_range {break}
-            assert!(cursor.goto_next_sibling());
-        }
-        //cursor.goto_parent();
-        // Broken due to https://github.com/tree-sitter/tree-sitter/issues/3270
-        //let _ = cursor.goto_first_child_for_byte(span_range.start).unwrap();
-        let start_node = cursor.node();
-        assert_eq!(start_node.kind_id(), kind);
-        assert_eq!(start_node.byte_range(), span_range);
-
-        Self{cursor, file_text, gathered_comments : Vec::new(), current_field_was_already_consumed : false}
-    }
-
     pub fn kind_span(&self) -> (u16, Span) {
         let node = self.cursor.node();
         (node.kind_id(), node.byte_range().into())