From b3d849ea7af6027ef7e9c9e2573abfe51d233f98 Mon Sep 17 00:00:00 2001
From: pca006132 <john.lck40@gmail.com>
Date: Wed, 21 Jul 2021 16:06:06 +0800
Subject: [PATCH] list test

---
 nac3core/src/typecheck/type_inferencer/mod.rs  | 6 ++++--
 nac3core/src/typecheck/type_inferencer/test.rs | 7 +++++++
 2 files changed, 11 insertions(+), 2 deletions(-)

diff --git a/nac3core/src/typecheck/type_inferencer/mod.rs b/nac3core/src/typecheck/type_inferencer/mod.rs
index 23b9660e3..624182b6d 100644
--- a/nac3core/src/typecheck/type_inferencer/mod.rs
+++ b/nac3core/src/typecheck/type_inferencer/mod.rs
@@ -304,7 +304,9 @@ impl<'a> Inferencer<'a> {
 
         Ok(Located {
             location,
-            custom: Some(list),
+            custom: Some(new_context.unifier.add_ty(TypeEnum::TList {
+                ty: elt.custom.unwrap(),
+            })),
             node: ExprKind::ListComp {
                 elt: Box::new(elt),
                 generators: vec![ast::Comprehension {
@@ -474,7 +476,7 @@ impl<'a> Inferencer<'a> {
         for t in elts.iter() {
             self.unifier.unify(ty, t.custom.unwrap())?;
         }
-        Ok(ty)
+        Ok(self.unifier.add_ty(TypeEnum::TList { ty }))
     }
 
     fn infer_tuple(&mut self, elts: &[ast::Expr<Option<Type>>]) -> InferenceResult {
diff --git a/nac3core/src/typecheck/type_inferencer/test.rs b/nac3core/src/typecheck/type_inferencer/test.rs
index c06208a3e..af9376232 100644
--- a/nac3core/src/typecheck/type_inferencer/test.rs
+++ b/nac3core/src/typecheck/type_inferencer/test.rs
@@ -159,6 +159,13 @@ impl TestEnvironment {
     [("a", "fn[[x=bool], bool]"), ("b", "fn[[x=int32], int32]"), ("c", "bool"),
      ("d", "int32"), ("foo1", "Foo[bool]"), ("foo2", "Foo[int32]")].iter().cloned().collect()
     ; "obj test")]
+#[test_case(indoc! {"
+        f = lambda x: True
+        a = [1, 2, 3]
+        b = [f(x) for x in a]
+    "},
+    [("a", "list[int32]"), ("b", "list[bool]"), ("f", "fn[[x=int32], bool]")].iter().cloned().collect()
+    ; "listcomp test")]
 fn test_basic(source: &str, mapping: HashMap<&str, &str>) {
     let mut env = TestEnvironment::new();
     let id_to_name = std::mem::take(&mut env.id_to_name);