crop return type

TODO.md

@@ -69,6 +69,7 @@
 - [x] Inner function
 - [x] Groover algorithm
 - [x] Tuple-tuple comparison
+- [x] Multi var assign
 
 ### Week 2: (30 Oct 23)
 

qlasskit/ast2ast.py

@@ -17,6 +17,21 @@
 from .ast2logic import flatten
 
 
+class NameConstantReplacer(ast.NodeTransformer):
+    def __init__(self, name_id, constant):
+        self.name_id = name_id
+        self.constant = constant
+
+    def generic_visit(self, node):
+        return super().generic_visit(node)
+
+    def visit_Name(self, node):
+        if node.id == self.name_id:
+            return ast.Constant(value=self.constant)
+
+        return node
+
+
 class ASTRewriter(ast.NodeTransformer):
     def __init__(self, env={}, ret=None):
         self.env = {}
@@ -145,7 +160,12 @@ def visit_For(self, node):
                 ast.Assign(targets=[node.target], value=ast.Constant(value=i))
             )
             rolls.extend(flatten([tar_assign]))
-            rolls.extend(flatten([self.visit(copy.deepcopy(b)) for b in node.body]))
+
+            new_body = [
+                NameConstantReplacer(node.target.id, i).visit(copy.deepcopy(b))
+                for b in node.body
+            ]
+            rolls.extend(flatten([self.visit(copy.deepcopy(b)) for b in new_body]))
 
         return rolls
 

qlasskit/ast2logic/t_expression.py

@@ -253,7 +253,7 @@ def unfold(v_exps, op):
         elif isinstance(expr.op, ast.BitXor) and hasattr(tleft[0], "bitwise_xor"):
             return tleft[0].bitwise_xor(tleft, tright)
         elif isinstance(expr.op, ast.BitAnd) and hasattr(tleft[0], "bitwise_and"):
-            return tleft[0].bitwise_and(tleft, tright)
+            return tright[0].bitwise_and(tleft, tright)  # type: ignore
         elif isinstance(expr.op, ast.BitOr) and hasattr(tleft[0], "bitwise_or"):
             return tleft[0].bitwise_or(tleft, tright)
         elif (

qlasskit/ast2logic/t_statement.py

@@ -71,6 +71,12 @@ def translate_statement(  # noqa: C901
             and texp.BIT_SIZE < ret_type.BIT_SIZE
         ):
             texp, vexp = ret_type.fill((texp, vexp))  # type: ignore
+        elif (
+            hasattr(texp, "BIT_SIZE")
+            and hasattr(ret_type, "BIT_SIZE")
+            and texp.BIT_SIZE > ret_type.BIT_SIZE
+        ):
+            texp, vexp = ret_type.crop((texp, vexp))  # type: ignore
         elif texp != ret_type:
             raise exceptions.TypeErrorException(texp, ret_type)
 

qlasskit/qlassf.py

@@ -284,6 +284,7 @@ def from_function(
 
         # Return the qlassf object
         qf = QlassF(fun_name, original_f, args, fun_ret, exps)
+
         if to_compile:
             qf.compile()
         return qf

qlasskit/types/qint.py

@@ -71,6 +71,16 @@ def fill(cls, v: TExp) -> TExp:
             )
         return v
 
+    @classmethod
+    def crop(cls, v: TExp) -> TExp:
+        """Crop a Qint to reach its bit_size"""
+        if len(v[1]) > cls.BIT_SIZE:  # type: ignore
+            v = (
+                cls,
+                v[1][0 : cls.BIT_SIZE],  # type: ignore
+            )
+        return v
+
     # Comparators
 
     @staticmethod

qlasskit/types/qtype.py

@@ -73,6 +73,11 @@ def fill(v: TExp) -> TExp:
         """Fill with leading false"""
         raise Exception("abstract")
 
+    @staticmethod
+    def crop(v: TExp) -> TExp:
+        """Crop to right size"""
+        raise Exception("abstract")
+
     # Comparators
 
     @staticmethod