From 56c845aac4548f16a15b6ebbeae48ab7e9aa1caa Mon Sep 17 00:00:00 2001
From: David Mak <chmakac@connect.ust.hk>
Date: Mon, 7 Oct 2024 17:00:20 +0800
Subject: [PATCH] [standalone] Add support for registering globals without type
 decl

---
 nac3core/src/toplevel/composer.rs | 13 +++++++
 nac3standalone/src/main.rs        | 62 +++++++++++++++----------------
 2 files changed, 42 insertions(+), 33 deletions(-)

diff --git a/nac3core/src/toplevel/composer.rs b/nac3core/src/toplevel/composer.rs
index f552229f..b93e1b36 100644
--- a/nac3core/src/toplevel/composer.rs
+++ b/nac3core/src/toplevel/composer.rs
@@ -382,6 +382,19 @@ impl TopLevelComposer {
                 ))
             }
 
+            ast::StmtKind::Assign { .. } => {
+                // Assignment statements can assign to (and therefore create) more than one 
+                // variable, but this function only allows returning one set of symbol information.
+                // We want to avoid changing this to return a `Vec` of symbol info, as this would 
+                // require `iter().next().unwrap()` on every variable created from a non-Assign 
+                // statement.
+                //
+                // Make callers use `register_top_level_var` instead, as it provides more
+                // fine-grained control over which symbols to register, while also simplifying the
+                // usage of this function.
+                panic!("Registration of top-level Assign statements must use TopLevelComposer::register_top_level_var (at {})", ast.location);
+            }
+
             ast::StmtKind::AnnAssign { target, annotation, .. } => {
                 let ExprKind::Name { id: name, .. } = target.node else {
                     return Err(format!(
diff --git a/nac3standalone/src/main.rs b/nac3standalone/src/main.rs
index 9b1a601c..71bfe8e5 100644
--- a/nac3standalone/src/main.rs
+++ b/nac3standalone/src/main.rs
@@ -174,46 +174,49 @@ fn handle_typevar_definition(
 fn handle_assignment_pattern(
     targets: &[Expr],
     value: &Expr,
-    resolver: &(dyn SymbolResolver + Send + Sync),
+    resolver: Arc<dyn SymbolResolver + Send + Sync>,
     internal_resolver: &ResolverInternal,
-    def_list: &[Arc<RwLock<TopLevelDef>>],
-    unifier: &mut Unifier,
-    primitives: &PrimitiveStore,
+    composer: &mut TopLevelComposer,
 ) -> Result<(), String> {
     if targets.len() == 1 {
-        match &targets[0].node {
+        let target = &targets[0];
+
+        match &target.node {
             ExprKind::Name { id, .. } => {
+                let def_list = composer.extract_def_list();
+                let unifier = &mut composer.unifier;
+                let primitives = &composer.primitives_ty;
+
                 if let Ok(var) =
-                    handle_typevar_definition(value, resolver, def_list, unifier, primitives)
+                    handle_typevar_definition(value, &*resolver, &def_list, unifier, primitives)
                 {
                     internal_resolver.add_id_type(*id, var);
                     Ok(())
-                } else if let Ok(val) = parse_parameter_default_value(value, resolver) {
+                } else if let Ok(val) = parse_parameter_default_value(value, &*resolver) {
                     internal_resolver.add_module_global(*id, val);
+                    let (name, def_id, _) = composer
+                        .register_top_level_var(
+                            *id,
+                            None,
+                            Some(resolver.clone()),
+                            "__main__",
+                            target.location,
+                        )
+                        .unwrap();
+                    internal_resolver.add_id_def(name, def_id);
                     Ok(())
                 } else {
                     Err(format!("fails to evaluate this expression `{:?}` as a constant or generic parameter at {}",
-                        targets[0].node,
-                        targets[0].location,
+                        target.node,
+                        target.location,
                     ))
                 }
             }
             ExprKind::List { elts, .. } | ExprKind::Tuple { elts, .. } => {
-                handle_assignment_pattern(
-                    elts,
-                    value,
-                    resolver,
-                    internal_resolver,
-                    def_list,
-                    unifier,
-                    primitives,
-                )?;
+                handle_assignment_pattern(elts, value, resolver, internal_resolver, composer)?;
                 Ok(())
             }
-            _ => Err(format!(
-                "assignment to {:?} is not supported at {}",
-                targets[0], targets[0].location
-            )),
+            _ => Err(format!("assignment to {target:?} is not supported at {}", target.location)),
         }
     } else {
         match &value.node {
@@ -223,11 +226,9 @@ fn handle_assignment_pattern(
                         handle_assignment_pattern(
                             std::slice::from_ref(tar),
                             val,
-                            resolver,
+                            resolver.clone(),
                             internal_resolver,
-                            def_list,
-                            unifier,
-                            primitives,
+                            composer,
                         )?;
                     }
                     Ok(())
@@ -360,17 +361,12 @@ fn main() {
     for stmt in parser_result {
         match &stmt.node {
             StmtKind::Assign { targets, value, .. } => {
-                let def_list = composer.extract_def_list();
-                let unifier = &mut composer.unifier;
-                let primitives = &composer.primitives_ty;
                 if let Err(err) = handle_assignment_pattern(
                     targets,
                     value,
-                    resolver.as_ref(),
+                    resolver.clone(),
                     internal_resolver.as_ref(),
-                    &def_list,
-                    unifier,
-                    primitives,
+                    &mut composer,
                 ) {
                     panic!("{err}");
                 }