From 3f4f23c0ab7ba6f355048f86743bef5037654ac4 Mon Sep 17 00:00:00 2001
From: pca006132 <jl@m-labs.hk>
Date: Wed, 23 Dec 2020 15:39:39 +0800
Subject: [PATCH] fix some variable inference

---
 toy-impl/parse_expr.py |  3 +++
 toy-impl/parse_stmt.py |  4 ++++
 toy-impl/top_level.py  | 13 ++++---------
 3 files changed, 11 insertions(+), 9 deletions(-)

diff --git a/toy-impl/parse_expr.py b/toy-impl/parse_expr.py
index aa0d945..f404440 100644
--- a/toy-impl/parse_expr.py
+++ b/toy-impl/parse_expr.py
@@ -264,6 +264,9 @@ def parse_list_comprehension(ctx: Context,
     if node.generators[0].is_async:
         raise CustomError('async list comprehension is not supported', node)
     ty = parse_expr(ctx, sym_table, node.generators[0].iter)
+    if isinstance(ty, TypeVariable) and \
+            len(ty.constraints) == 1:
+        ty = ty.constraints[0]
     if not isinstance(ty, ListType):
         raise CustomError(f'unable to iterate over {ty}', node)
     try:
diff --git a/toy-impl/parse_stmt.py b/toy-impl/parse_stmt.py
index 075b083..6c60f05 100644
--- a/toy-impl/parse_stmt.py
+++ b/toy-impl/parse_stmt.py
@@ -159,7 +159,11 @@ def parse_for_stmt(ctx: Context,
                    return_ty: Type,
                    node):
     ty = parse_expr(ctx, sym_table, node.iter)
+    if isinstance(ty, TypeVariable) and \
+            len(ty.constraints) == 1:
+        ty = ty.constraints[0]
     if not isinstance(ty, ListType):
+
         raise CustomError('only iteration over list is supported', node.iter)
     binding = parse_simple_binding(node.target, ty.params[0])
     for key, value in binding.items():
diff --git a/toy-impl/top_level.py b/toy-impl/top_level.py
index 86594c7..fbd18b5 100644
--- a/toy-impl/top_level.py
+++ b/toy-impl/top_level.py
@@ -126,17 +126,12 @@ def parse_type_var(ctx: Context, node):
             raise CustomError('redefining type variable is not allowed', node)
         constraints = []
         for v in node.value.args[1:]:
-            if isinstance(v, ast.Constant):
-                value = v.value
-            elif isinstance(v, ast.Name):
-                value = v.id
-            else:
+            value, var = parse_type(ctx, v)
+            if len(var) > 0:
                 raise CustomError(
-                    'TypeVar constraints must be either string or type name',
+                    'constraints cannot contain type variables',
                     node)
-            if value not in ctx.types:
-                raise CustomError(f'unbounded type {value}', node)
-            constraints.append(ctx.types[value])
+            constraints.append(value)
         ctx.variables[name] = TypeVariable(node.value.args[0].value, constraints)