From 1407553a2f118ba4d95f4d81cb2f74e0ff7b79b4 Mon Sep 17 00:00:00 2001
From: David Mak <chmakac@connect.ust.hk>
Date: Fri, 4 Oct 2024 13:08:09 +0800
Subject: [PATCH] [core] Implement parsing of global variables

Globals are now parsed into symbol resolver and top level definitions.
---
 nac3core/src/toplevel/composer.rs | 54 +++++++++++++++++++++++++++++--
 nac3standalone/src/main.rs        | 45 ++++++++++++++++++++++++++
 2 files changed, 97 insertions(+), 2 deletions(-)

diff --git a/nac3core/src/toplevel/composer.rs b/nac3core/src/toplevel/composer.rs
index 5906e159..06d5022b 100644
--- a/nac3core/src/toplevel/composer.rs
+++ b/nac3core/src/toplevel/composer.rs
@@ -1,6 +1,6 @@
 use std::rc::Rc;
 
-use nac3parser::ast::fold::Fold;
+use nac3parser::ast::{fold::Fold, ExprKind};
 
 use super::*;
 use crate::{
@@ -382,8 +382,58 @@ impl TopLevelComposer {
                 ))
             }
 
+            ast::StmtKind::AnnAssign { target, annotation, .. } => {
+                let ExprKind::Name { id: name, .. } = target.node else {
+                    return Err(format!(
+                        "global variable declaration must be an identifier (at {})",
+                        ast.location
+                    ));
+                };
+
+                if self.keyword_list.contains(&name) {
+                    return Err(format!(
+                        "cannot use keyword `{}` as a class name (at {})",
+                        name,
+                        ast.location
+                    ));
+                }
+
+                let global_var_name = if mod_path.is_empty() {
+                    name.to_string()
+                } else {
+                    format!("{mod_path}.{name}")
+                };
+                if !defined_names.insert(global_var_name.clone()) {
+                    return Err(format!(
+                        "global variable `{}` defined twice (at {})",
+                        global_var_name,
+                        ast.location
+                    ));
+                }
+
+                let ty_to_be_unified = self.unifier.get_dummy_var().ty;
+                self.definition_ast_list.push((
+                    RwLock::new(Self::make_top_level_variable_def(
+                        global_var_name,
+                        name,
+                        // dummy here, unify with correct type later,
+                        ty_to_be_unified,
+                        *(annotation.clone()),
+                        resolver,
+                        Some(ast.location),
+                    )).into(),
+                    None,
+                ));
+
+                Ok((
+                    name,
+                    DefinitionId(self.definition_ast_list.len() - 1),
+                    Some(ty_to_be_unified),
+                ))
+            }
+
             _ => Err(format!(
-                "registrations of constructs other than top level classes/functions are not supported (at {})",
+                "registrations of constructs other than top level classes/functions/variables are not supported (at {})",
                 ast.location
             )),
         }
diff --git a/nac3standalone/src/main.rs b/nac3standalone/src/main.rs
index 02902187..e27c5e1f 100644
--- a/nac3standalone/src/main.rs
+++ b/nac3standalone/src/main.rs
@@ -245,6 +245,34 @@ fn handle_assignment_pattern(
     }
 }
 
+fn handle_global_var(
+    target: &Expr,
+    value: Option<&Expr>,
+    resolver: &(dyn SymbolResolver + Send + Sync),
+    internal_resolver: &ResolverInternal,
+) -> Result<(), String> {
+    let ExprKind::Name { id, .. } = target.node else {
+        return Err(format!(
+            "global variable declaration must be an identifier (at {})",
+            target.location,
+        ));
+    };
+
+    let Some(value) = value else {
+        return Err(format!("global variable `{id}` must be initialized in its definition"));
+    };
+
+    if let Ok(val) = parse_parameter_default_value(value, resolver) {
+        internal_resolver.add_module_global(id, val);
+        Ok(())
+    } else {
+        Err(format!(
+            "failed to evaluate this expression `{:?}` as a constant at {}",
+            target.node, target.location,
+        ))
+    }
+}
+
 fn main() {
     let cli = CommandLineArgs::parse();
     let CommandLineArgs { file_name, threads, opt_level, emit_llvm, triple, mcpu, target_features } =
@@ -342,6 +370,23 @@ fn main() {
                     panic!("{err}");
                 }
             }
+
+            StmtKind::AnnAssign { target, value, .. } => {
+                if let Err(err) = handle_global_var(
+                    target,
+                    value.as_ref().map(Box::as_ref),
+                    resolver.as_ref(),
+                    internal_resolver.as_ref(),
+                ) {
+                    panic!("{err}");
+                }
+
+                let (name, def_id, _) = composer
+                    .register_top_level(stmt, Some(resolver.clone()), "__main__", true)
+                    .unwrap();
+                internal_resolver.add_id_def(name, def_id);
+            }
+
             // allow (and ignore) "from __future__ import annotations"
             StmtKind::ImportFrom { module, names, .. }
                 if module == &Some("__future__".into())