add return type to qlassf

dakk · Oct 19, 2023 · 023896d · 023896d
1 parent c4ab83b
commit 023896d
Show file tree

Hide file tree

Showing 3 changed files with 13 additions and 9 deletions.
diff --git a/qlasskit/ast2logic/t_ast.py b/qlasskit/ast2logic/t_ast.py
@@ -38,8 +38,7 @@ def translate_ast(fun, types) -> LogicFun:
     if not fun.returns:
         raise exceptions.NoReturnTypeException()
 
-    ret_ = translate_argument(fun.returns, env)  # TODO: we need to preserve this
-    ret_size = len(ret_)
+    ret_ = translate_argument(fun.returns, env)
 
     exps = []
     for stmt in fun.body:
@@ -53,4 +52,4 @@ def translate_ast(fun, types) -> LogicFun:
     #     if e == true or e == false:
     #         print(f"Warning: expression {n} is returning a costant: {e}")
 
-    return fun_name, args, ret_size, exps_simpl
+    return fun_name, args, ret_, exps_simpl
diff --git a/qlasskit/ast2logic/typing.py b/qlasskit/ast2logic/typing.py
@@ -38,4 +38,4 @@ def to_exp(self) -> List[Symbol]:
 
 Args = List[Arg]
 BoolExpList = List[Tuple[Symbol, Boolean]]
-LogicFun = Tuple[str, Args, int, BoolExpList]
+LogicFun = Tuple[str, Args, Arg, BoolExpList]
diff --git a/qlasskit/qlassf.py b/qlasskit/qlassf.py
@@ -19,7 +19,7 @@
 
 from . import compiler
 from .ast2ast import ast2ast
-from .ast2logic import Args, BoolExpList, flatten, translate_ast
+from .ast2logic import Arg, Args, BoolExpList, flatten, translate_ast
 from .types import *  # noqa: F403, F401
 from .types import Qtype
 
@@ -32,29 +32,34 @@ class QlassF:
     name: str
     original_f: Callable
     args: Args
-    ret_size: int
+    returns: Arg
     expressions: BoolExpList
 
     def __init__(
         self,
         name: str,
         original_f: Callable,
         args: Args,
-        ret_size: int,
+        returns: Arg,
         exps: BoolExpList,
     ):
         self.name = name
         self.original_f = original_f
         self.args = args
-        self.ret_size = ret_size
+        self.returns = returns
         self.expressions = exps
 
         self._compiled_gate = None
 
     def __repr__(self):
+        ret_str = f"{self.returns.ttype}[{len(self.returns)}]"
         arg_str = ", ".join(map(lambda arg: f"{arg.name}:{arg.ttype}", self.args))
         exp_str = "\n\t".join(map(lambda exp: f"{exp[0]} = {exp[1]}", self.expressions))
-        return f"QlassF<{self.name}>({arg_str}) -> bool[{self.ret_size}]:\n\t{exp_str}"
+        return f"QlassF<{self.name}>({arg_str}) -> {ret_str}:\n\t{exp_str}"
+
+    @property
+    def ret_size(self):
+        return len(self.returns)
 
     def __add__(self, qf2) -> "QlassF":
         """Adds two qlassf and return the combination"""