From e8a5843ca70bfd2788fd8f324d0e3be0efe9c9b3 Mon Sep 17 00:00:00 2001
From: ychenfo <yc@m-labs.hk>
Date: Thu, 11 Nov 2021 01:54:29 +0800
Subject: [PATCH] nac3standalone: basic resolver typevar handling

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

diff --git a/nac3core/src/toplevel/composer.rs b/nac3core/src/toplevel/composer.rs
index bf9f1561..013829e0 100644
--- a/nac3core/src/toplevel/composer.rs
+++ b/nac3core/src/toplevel/composer.rs
@@ -534,7 +534,7 @@ impl TopLevelComposer {
         }
     }
 
-    fn extract_def_list(&self) -> Vec<Arc<RwLock<TopLevelDef>>> {
+    pub fn extract_def_list(&self) -> Vec<Arc<RwLock<TopLevelDef>>> {
         self.definition_ast_list.iter().map(|(def, ..)| def.clone()).collect_vec()
     }
 
diff --git a/nac3standalone/src/main.rs b/nac3standalone/src/main.rs
index 5b4e31dd..a4a82e53 100644
--- a/nac3standalone/src/main.rs
+++ b/nac3standalone/src/main.rs
@@ -4,7 +4,7 @@ use inkwell::{
     OptimizationLevel,
 };
 use nac3core::typecheck::type_inferencer::PrimitiveStore;
-use nac3parser::parser;
+use nac3parser::{ast::{ExprKind, StmtKind}, parser};
 use std::env;
 use std::fs;
 use std::{collections::HashMap, path::Path, sync::Arc, time::SystemTime};
@@ -66,6 +66,46 @@ fn main() {
     );
 
     for stmt in parser_result.into_iter() {
+        // handle type vars in toplevel
+        if let StmtKind::Assign { value, targets, .. } = &stmt.node {
+            assert_eq!(targets.len(), 1, "only support single assignment for now, at {}", targets[0].location);
+            if let ExprKind::Call { func, args, .. } = &value.node {
+                if matches!(&func.node, ExprKind::Name { id, .. } if id == &"TypeVar".into()) {
+                    print!("registering typevar {:?}", targets[0].node);
+                    let constraints = args
+                        .iter()
+                        .skip(1)
+                        .map(|x| {
+                            let def_list = &composer.extract_def_list();
+                            let unifier = &mut composer.unifier;
+                            resolver.parse_type_annotation(
+                                def_list,
+                                unifier,
+                                &primitive,
+                                x
+                            ).unwrap()
+                        })
+                        .collect::<Vec<_>>();
+                    let res_ty = composer.unifier.get_fresh_var_with_range(&constraints).0;
+                    println!(
+                        " ...registered: {}",
+                        composer.unifier.stringify(
+                            res_ty,
+                            &mut |x| format!("obj{}", x),
+                            &mut |x| format!("tavr{}", x)
+                        )
+                    );
+                    internal_resolver.add_id_type(
+                        if let ExprKind::Name { id, .. } = &targets[0].node { *id } else {
+                            panic!("must assign simple name variable as type variable for now")
+                        },
+                        res_ty
+                    );
+                    continue;
+                }
+            }
+        }
+
         let (name, def_id, ty) = composer
             .register_top_level(stmt, Some(resolver.clone()), "__main__".into())
             .unwrap();