From e60e8e837f43572b5997124488e3c8a7e680c5b4 Mon Sep 17 00:00:00 2001
From: David Mak <chmakac@connect.ust.hk>
Date: Fri, 4 Oct 2024 13:15:00 +0800
Subject: [PATCH] [core] Add support for global statements

---
 nac3core/src/codegen/stmt.rs                  | 31 +++++++++++++++++
 nac3core/src/typecheck/function_check.rs      | 34 +++++++++++++++++++
 nac3core/src/typecheck/type_inferencer/mod.rs |  1 +
 3 files changed, 66 insertions(+)

diff --git a/nac3core/src/codegen/stmt.rs b/nac3core/src/codegen/stmt.rs
index fb81c6dc2..cfc188c87 100644
--- a/nac3core/src/codegen/stmt.rs
+++ b/nac3core/src/codegen/stmt.rs
@@ -1828,6 +1828,37 @@ pub fn gen_stmt<G: CodeGenerator>(
                 stmt.location,
             );
         }
+        StmtKind::Global { names, .. } => {
+            let registered_globals = ctx
+                .top_level
+                .definitions
+                .read()
+                .iter()
+                .filter_map(|def| {
+                    if let TopLevelDef::Variable { simple_name, ty, .. } = &*def.read() {
+                        Some((*simple_name, *ty))
+                    } else {
+                        None
+                    }
+                })
+                .collect_vec();
+
+            for id in names {
+                let Some((_, ty)) = registered_globals.iter().find(|(name, _)| name == id) else {
+                    return Err(format!("{id} is not a global at {}", stmt.location));
+                };
+
+                let resolver = ctx.resolver.clone();
+                let ptr = resolver
+                    .get_symbol_value(*id, ctx, generator)
+                    .map(|val| val.to_basic_value_enum(ctx, generator, *ty))
+                    .transpose()?
+                    .map(BasicValueEnum::into_pointer_value)
+                    .unwrap();
+
+                ctx.var_assignment.insert(*id, (ptr, None, 0));
+            }
+        }
         _ => unimplemented!(),
     };
     Ok(())
diff --git a/nac3core/src/typecheck/function_check.rs b/nac3core/src/typecheck/function_check.rs
index 4b5ccb579..0f38a341f 100644
--- a/nac3core/src/typecheck/function_check.rs
+++ b/nac3core/src/typecheck/function_check.rs
@@ -365,6 +365,40 @@ impl<'a> Inferencer<'a> {
                 }
                 Ok(true)
             }
+            StmtKind::Global { names, .. } => {
+                for id in names {
+                    if let Some(id_info) = defined_identifiers.get(id) {
+                        if !id_info.is_global {
+                            return Err(HashSet::from([format!(
+                                "name '{id}' is assigned to before global declaration at {}",
+                                stmt.location,
+                            )]));
+                        }
+
+                        continue;
+                    }
+
+                    match self.function_data.resolver.get_symbol_type(
+                        self.unifier,
+                        &self.top_level.definitions.read(),
+                        self.primitives,
+                        *id,
+                    ) {
+                        Ok(_) => {
+                            self.defined_identifiers
+                                .insert(*id, IdentifierInfo { is_global: true });
+                        }
+                        Err(e) => {
+                            return Err(HashSet::from([format!(
+                                "type error at identifier `{}` ({}) at {}",
+                                id, e, stmt.location
+                            )]))
+                        }
+                    }
+                }
+
+                Ok(false)
+            }
             // break, raise, etc.
             _ => Ok(false),
         }
diff --git a/nac3core/src/typecheck/type_inferencer/mod.rs b/nac3core/src/typecheck/type_inferencer/mod.rs
index ba09bd5aa..ce2e76d27 100644
--- a/nac3core/src/typecheck/type_inferencer/mod.rs
+++ b/nac3core/src/typecheck/type_inferencer/mod.rs
@@ -394,6 +394,7 @@ impl<'a> Fold<()> for Inferencer<'a> {
             | ast::StmtKind::Continue { .. }
             | ast::StmtKind::Expr { .. }
             | ast::StmtKind::For { .. }
+            | ast::StmtKind::Global { .. }
             | ast::StmtKind::Pass { .. }
             | ast::StmtKind::Try { .. } => {}
             ast::StmtKind::If { test, .. } | ast::StmtKind::While { test, .. } => {