Add support for assignment expressions

weirdsmiley · weirdsmiley · commit b3871b20f3b1 · 2025-12-16T15:29:18.000+05:30
We now have two kinds of expressions, let and non-let. They mimick
def-use chains. Also, I have mandated the usage of `return` keyword.
diff --git a/examples/math.ql b/examples/math.ql
@@ -8,6 +8,6 @@ fn cos(x: f64) : f64 {
 
 fn exp(x: f64) : f64 {
     let e = 2.718;
-    e * x
+    return e * x;
 }
 
diff --git a/examples/toss.ql b/examples/toss.ql
@@ -3,15 +3,15 @@ import std::Hadamard;
 fn toss() : qbit {
     let zero_state = 0q(0, 1);  // represent a qubit in zero state simply as 0
     let superpositioned = Hadamard(zero_state);
-    superpositioned
+    return superpositioned;
 }
 
 fn main() {
     let choice = toss();
     if choice == 0 {
-        print("Heads");
+        // print("Heads");
     } else {
-        print("Tails");
+        // print("Tails");
     }
 }
 
diff --git a/src/ast.rs b/src/ast.rs
@@ -590,6 +590,7 @@ pub enum Expr {
     Tensor(Vec<QccCell<Expr>>),
     FnCall(FunctionAST, Vec<QccCell<Expr>>),
     Let(VarAST, QccCell<Expr>),
+    Assign(VarAST, QccCell<Expr>),
     Conditional(
         QccCell<Expr>,
         Vec<QccCell<Expr>>, /* truth_block */
@@ -606,6 +607,7 @@ impl Expr {
             Self::Tensor(_) => Default::default(), // TODO: This will require subtracting the dimension of tensor
             Self::FnCall(f, _) => f.get_loc().clone(),
             Self::Let(var, _) => var.location.clone(),
+            Self::Assign(v, _) => v.location().clone(),
             Self::Conditional(c, _, _) => c.as_ref().borrow().get_location(),
             Self::Literal(lit) =>
             /*TODO*/
@@ -643,6 +645,7 @@ impl Expr {
             }
             Self::FnCall(f, args) => *f.get_output_type(),
             Self::Let(var, val) => var.get_type(),
+            Self::Assign(var, val) => var.get_type(),
             Self::Conditional(conditional, truth_block, false_block) => {
                 let last_expr = truth_block.last();
                 if last_expr.is_none() {
@@ -709,6 +712,7 @@ impl std::fmt::Display for Expr {
 
             Self::Let(var, val) => write!(f, "{} = {}", var, *val.as_ref().borrow()),
 
+            Self::Assign(var, val) => write!(f, "{} = {}", var, *val.as_ref().borrow()),
             Self::Conditional(cond, _truth, _false) => {
                 writeln!(f, "{}", *cond.as_ref().borrow())?;
 
diff --git a/src/inference.rs b/src/inference.rs
@@ -119,6 +119,18 @@ fn check_expr(expr: &QccCell<Expr>) -> Result<Type> {
 
             Ok(val_type)
         }
+        Expr::Assign(ref var, ref val) => {
+            if !var.is_typed() {
+                return Err(QccErrorKind::UnknownType)?;
+            }
+            let val_type = check_expr(val)?;
+
+            if var.get_type() != val_type {
+                return Err(QccErrorKind::TypeMismatch)?;
+            }
+
+            Ok(val_type)
+        }
         Expr::Conditional(ref conditional, ref truth_block, ref false_block) => {
             for expr in truth_block {
                 check_expr(expr);
@@ -435,6 +447,23 @@ fn infer_expr(expr: &QccCell<Expr>) -> Option<Type> {
             }
         }
 
+        Expr::Assign(ref mut var, ref val) => {
+            // val is an expression and it must have the same type as var
+            if var.get_type() == Type::Bottom {
+                // we need to type check from expression first
+                let rhs_type = infer_expr(&val)?;
+                var.set_type(rhs_type);
+                return Some(rhs_type);
+            } else {
+                let lhs_type = var.get_type();
+                let rhs_type = infer_expr(&val)?;
+                if lhs_type != rhs_type {
+                    return None;
+                }
+                return Some(lhs_type);
+            }
+        }
+
         Expr::Conditional(ref conditional, ref truth_block, ref false_block) => {
             let mut truth_block_type = Some(Type::Bottom);
             for expr in truth_block {
@@ -646,6 +675,55 @@ fn infer_from_table(
                 .into()))
             }
         }
+        Expr::Assign(ref mut var, ref val) => {
+            let rhs_info = infer_from_table(val, param_st, local_st, function_st);
+
+            if rhs_info.is_some() {
+                return rhs_info;
+            }
+
+            let var_type = var.get_type();
+            let val_type = val.as_ref().borrow().get_type();
+
+            if !var.is_typed() {
+                var.set_type(val.as_ref().borrow().get_type());
+                None
+            } else if (var_type == Type::Qbit || var_type == Type::Bit)
+                && (val_type == Type::Qbit || val_type == Type::Bit)
+                && (var_type != val_type)
+            {
+                // If one of lhs or rhs is qbit while the other is bit, then we
+                // will put a measure operator before it is assigned during
+                // codegen.
+                //
+                //  let x: bit = y;     # y := qbit
+                //
+                // This holds according to our subtyping rules. Codegen will
+                // lower this to:
+                //
+                //  let x0: bit = measure(y);
+                //  let x: bit  = x0;
+                //
+                // Similarily,
+                //
+                //  let x: qbit = y;    # y := bit
+                //
+                // This is also fine. When codegen lowers this code, it
+                // automatically puts required stub to create a logical qubit.
+                None
+            } else if var_type == val_type {
+                None
+            } else if var_type != val_type {
+                // if one is qbit and other is bit, pass
+                Some(Err(QccErrorKind::TypeMismatch.into()))
+            } else {
+                Some(Ok(Expr::Var(VarAST::new(
+                    var.name().clone(),
+                    var.location().clone(),
+                ))
+                .into()))
+            }
+        }
         Expr::Conditional(ref mut conditional, ref mut truth_block, ref mut false_block) => {
             for expr in truth_block {
                 let info = infer_from_table(expr, param_st, local_st, function_st);
diff --git a/src/mangle.rs b/src/mangle.rs
@@ -31,6 +31,9 @@ pub(crate) fn mangle_expr(expr: &mut QccCell<Expr>, prefix: Ident) {
         Expr::Let(_, ref mut val) => {
             mangle_expr(val, prefix);
         }
+        Expr::Assign(_, ref mut val) => {
+            mangle_expr(val, prefix);
+        }
         Expr::FnCall(ref mut f, ref mut args) => {
             for arg in args {
                 mangle_expr(arg, prefix.clone());
@@ -66,6 +69,9 @@ pub(crate) fn mangle_fns(expr: &mut QccCell<Expr>, module_name: &String, functio
         Expr::Let(ref mut var, ref mut val) => {
             mangle_fns(val, module_name, functions);
         }
+        Expr::Assign(ref mut var, ref mut val) => {
+            mangle_fns(val, module_name, functions);
+        }
         Expr::FnCall(ref mut f, ref mut args) => {
             for arg in args {
                 mangle_fns(arg, module_name, functions);
@@ -109,6 +115,9 @@ fn mangle_expr_check(expr: &mut QccCell<Expr>, mod_name: &Ident, fn_name: &Ident
         Expr::Let(_, ref mut val) => {
             mangle_expr_check(val, mod_name, fn_name);
         }
+        Expr::Assign(_, ref mut val) => {
+            mangle_expr_check(val, mod_name, fn_name);
+        }
         Expr::FnCall(ref mut f, ref mut args) => {
             for arg in args {
                 mangle_expr_check(arg, mod_name, fn_name);
diff --git a/src/parser.rs b/src/parser.rs
@@ -194,6 +194,22 @@ impl Parser {
             } else if self.lexer.is_token(Token::If) {
                 let expr = self.parse_if_block()?;
                 body.push(expr);
+            } else if self.lexer.is_token(Token::Identifier) {
+                // Normal assignment expressions
+                let name = self.lexer.identifier();
+                let location = self.lexer.location.clone();
+                self.lexer.consume(Token::Identifier)?;
+
+                let var = VarAST::new(name.clone(), location.clone());
+
+                if !self.lexer.is_token(Token::Assign) {
+                    return Err(QccErrorKind::ExpectedAssign)?;
+                }
+                self.lexer.consume(Token::Assign)?;
+
+                let val = self.parse_expr()?;
+                let expr = Expr::Assign(var, val).into();
+                body.push(expr);
             } else {
                 // function calls with no return or dropped return value.
                 if self.lexer.token.is_some() {
diff --git a/tests/let_both_typed.ql b/tests/let_both_typed.ql
@@ -1,14 +1,14 @@
 fn foo() : qbit {
     let q = 0q(0, 1);  // represent a qubit in zero state simply as 0
-    q
+    return q;
 }
 
 fn main() {
     let choice : qbit = foo();
     if choice == 0 {
-        print("Heads");
+        let _ = print(0);
     } else {
-        print("Tails");
+        let _ = print(1);
     }
 }
 
diff --git a/tests/package/math.ql b/tests/package/math.ql
@@ -10,6 +10,6 @@ fn cos(x: f64) : f64 {
 
 fn exp(x: f64) : f64 {
     let e = 2.718;
-    e * x
+    return e * x;
 }
 
diff --git a/tests/package/toss.ql b/tests/package/toss.ql
@@ -3,15 +3,15 @@ import std::Hadamard;
 fn toss() : qbit {
     let zero_state = 0q(0, 1);  // represent a qubit in zero state simply as 0
     let superpositioned = Hadamard(zero_state);
-    superpositioned
+    return superpositioned;
 }
 
 fn main() {
     let choice = toss();
     if choice == 0 {
-        print("Heads");
+        // print("Heads");
     } else {
-        print("Tails");
+        // print("Tails");
     }
 }
 
diff --git a/tests/qbit_float.ql b/tests/qbit_float.ql
@@ -1,11 +1,11 @@
 fn foo(q0: qbit) : qbit {
     let q1 = 2 * q0;
-    q1
+    return q1;
 }
 
 fn bar(q0: qbit) : qbit {
     let q1 = q0 * 2;
-    q1
+    return q1;
 }
 
 fn main() {
diff --git a/tests/tensors2.ql b/tests/tensors2.ql
@@ -1,17 +1,17 @@
-fn sin(x: f64) : f64 {
-    x
+fn foo(x: f64) : f64 {
+    return x;
 }
 
-fn cos(x: f64) : f64 {
-    x
+fn bar(x: f64) : f64 {
+    return x;
 }
 
 fn main() {
     let t1 = [[], []];
     let t2 = [[[], []], [[]]];
 }
 
-fn foo() {
+fn foo2() {
     let x = 42;
     let e0 = 2.718;
     let e1 = e0 * 2;
@@ -20,10 +20,10 @@ fn foo() {
     let t1 = [];
     let t2 = [x];
     let t3 = [t2, t2];
-    let t4 = [e0 * cos(a), -e1 * sin(a)];
+    let t4 = [e0 * bar(a), -e1 * foo(a)];
     let t5 = [[]];
     let t7 = [[x]];
-    let t8 = [[e0 * cos(a), -e1 * sin(a)], [e1 * sin(a),  e0 * cos(a)]];
+    let t8 = [[e0 * bar(a), -e1 * foo(a)], [e1 * foo(a),  e0 * bar(a)]];
     let t6 = [[], []]; // FIXME
     let t9 = [[[], []], [[]]]; // FIXME
     return t3;
diff --git a/tests/tests.rs b/tests/tests.rs
@@ -99,9 +99,11 @@ fn test_ast_gen() -> Result<(), Box<dyn std::error::Error>>  {
 "|_ qbit_float			// @qbit_float.ql:1:1
   |_ fn qbit_float$foo (q0: qubit) : qubit		// @qbit_float.ql:1:4
     |_ q1: qubit = (2 * q0: qubit)
+    |_ q1: qubit
 
   |_ fn qbit_float$bar (q0: qubit) : qubit		// @qbit_float.ql:6:4
     |_ q1: qubit = (q0: qubit * 2)
+    |_ q1: qubit
 
   |_ fn qbit_float$main () : qubit		// @qbit_float.ql:11:4
     |_ x: qubit = qbit_float$foo: qubit ()
@@ -120,26 +122,28 @@ fn test_ast_gen() -> Result<(), Box<dyn std::error::Error>>  {
 
     test!("tests/tensors2.ql",
 "|_ tensors2			// @tensors2.ql:1:1
-  |_ fn tensors2$sin (x: float64) : float64		// @tensors2.ql:1:4
+  |_ fn tensors2$foo (x: float64) : float64		// @tensors2.ql:1:4
+    |_ x: float64
 
-  |_ fn tensors2$cos (x: float64) : float64		// @tensors2.ql:5:4
+  |_ fn tensors2$bar (x: float64) : float64		// @tensors2.ql:5:4
+    |_ x: float64
 
   |_ fn tensors2$main () : <bottom>		// @tensors2.ql:9:4
     |_ t1 = [[], []]
     |_ t2 = [[[], []], [[]]]
 
-  |_ fn tensors2$foo () : float64		// @tensors2.ql:14:4
+  |_ fn tensors2$foo2 () : float64		// @tensors2.ql:14:4
     |_ x: float64 = 42
     |_ e0: float64 = 2.718
     |_ e1: float64 = (e0: float64 * 2)
     |_ a: float64 = 0.707
     |_ t1 = []
     |_ t2: float64 = [x: float64]
     |_ t3: float64 = [t2: float64, t2: float64]
-    |_ t4: float64 = [(e0: float64 * tensors2$cos: float64 (a: float64)), (-e1: float64 * tensors2$sin: float64 (a: float64))]
+    |_ t4: float64 = [(e0: float64 * tensors2$bar: float64 (a: float64)), (-e1: float64 * tensors2$foo: float64 (a: float64))]
     |_ t5 = [[]]
     |_ t7: float64 = [[x: float64]]
-    |_ t8: float64 = [[(e0: float64 * tensors2$cos: float64 (a: float64)), (-e1: float64 * tensors2$sin: float64 (a: float64))], [(e1: float64 * tensors2$sin: float64 (a: float64)), (e0: float64 * tensors2$cos: float64 (a: float64))]]
+    |_ t8: float64 = [[(e0: float64 * tensors2$bar: float64 (a: float64)), (-e1: float64 * tensors2$foo: float64 (a: float64))], [(e1: float64 * tensors2$foo: float64 (a: float64)), (e0: float64 * tensors2$bar: float64 (a: float64))]]
     |_ t6 = [[], []]
     |_ t9 = [[[], []], [[]]]
     |_ t3: float64
@@ -305,6 +309,7 @@ fn test_ast_gen() -> Result<(), Box<dyn std::error::Error>>  {
   |_ fn toss$toss () : qubit		// @toss.ql:3:4
     |_ zero_state: qubit = 0q0_1
     |_ superpositioned: qubit = std$Hadamard: qubit (zero_state: qubit)
+    |_ superpositioned: qubit
 
   |_ fn toss$main () : <bottom>		// @toss.ql:9:4
     |_ choice: qubit = toss$toss: qubit ()
@@ -332,6 +337,7 @@ fn test_ast_gen() -> Result<(), Box<dyn std::error::Error>>  {
 
   |_ fn math$exp (x: float64) : float64		// @math.ql:9:4
     |_ e: float64 = 2.718
+    |_ (e: float64 * x: float64)
 
 ");
     Ok(())

Original file line number	Diff line number	Diff line change
`@@ -8,6 +8,6 @@ fn cos(x: f64) : f64 {`
`8`	`8`
`9`	`9`	`fn exp(x: f64) : f64 {`
`10`	`10`	`let e = 2.718;`
`11`		`- e * x`
	`11`	`+ return e * x;`
`12`	`12`	`}`
`13`	`13`
Original file line number	Diff line number	Diff line change
`@@ -3,15 +3,15 @@ import std::Hadamard;`
`3`	`3`	`fn toss() : qbit {`
`4`	`4`	`let zero_state = 0q(0, 1); // represent a qubit in zero state simply as 0`
`5`	`5`	`let superpositioned = Hadamard(zero_state);`
`6`		`- superpositioned`
	`6`	`+ return superpositioned;`
`7`	`7`	`}`
`8`	`8`
`9`	`9`	`fn main() {`
`10`	`10`	`let choice = toss();`
`11`	`11`	`if choice == 0 {`
`12`		`- print("Heads");`
	`12`	`+ // print("Heads");`
`13`	`13`	`} else {`
`14`		`- print("Tails");`
	`14`	`+ // print("Tails");`
`15`	`15`	`}`
`16`	`16`	`}`
`17`	`17`
Original file line number	Diff line number	Diff line change
`@@ -1,14 +1,14 @@`
`1`	`1`	`fn foo() : qbit {`
`2`	`2`	`let q = 0q(0, 1); // represent a qubit in zero state simply as 0`
`3`		`- q`
	`3`	`+ return q;`
`4`	`4`	`}`
`5`	`5`
`6`	`6`	`fn main() {`
`7`	`7`	`let choice : qbit = foo();`
`8`	`8`	`if choice == 0 {`
`9`		`- print("Heads");`
	`9`	`+ let _ = print(0);`
`10`	`10`	`} else {`
`11`		`- print("Tails");`
	`11`	`+ let _ = print(1);`
`12`	`12`	`}`
`13`	`13`	`}`
`14`	`14`
Original file line number	Diff line number	Diff line change
`@@ -10,6 +10,6 @@ fn cos(x: f64) : f64 {`
`10`	`10`
`11`	`11`	`fn exp(x: f64) : f64 {`
`12`	`12`	`let e = 2.718;`
`13`		`- e * x`
	`13`	`+ return e * x;`
`14`	`14`	`}`
`15`	`15`
Original file line number	Diff line number	Diff line change
`@@ -1,11 +1,11 @@`
`1`	`1`	`fn foo(q0: qbit) : qbit {`
`2`	`2`	`let q1 = 2 * q0;`
`3`		`- q1`
	`3`	`+ return q1;`
`4`	`4`	`}`
`5`	`5`
`6`	`6`	`fn bar(q0: qbit) : qbit {`
`7`	`7`	`let q1 = q0 * 2;`
`8`		`- q1`
	`8`	`+ return q1;`
`9`	`9`	`}`
`10`	`10`
`11`	`11`	`fn main() {`