More comprehensive update for variable updates in static exprs

source/compiler/qsc_eval/src/lib.rs

@@ -453,7 +453,7 @@ impl Env {
             .find_map(|scope| scope.bindings.get(id))
     }
 
-    fn get_mut(&mut self, id: LocalVarId) -> Option<&mut Variable> {
+    pub fn get_mut(&mut self, id: LocalVarId) -> Option<&mut Variable> {
         self.scopes
             .iter_mut()
             .rev()

source/compiler/qsc_partial_eval/src/evaluation_context.rs

@@ -9,10 +9,10 @@ use qsc_eval::{
 use qsc_fir::fir::{LocalItemId, LocalVarId, PackageId};
 use qsc_rca::{ComputeKind, RuntimeFeatureFlags, ValueKind};
 use qsc_rir::rir::{BlockId, Literal, VariableId};
-use rustc_hash::FxHashMap;
+use rustc_hash::{FxHashMap, FxHashSet};
 use std::{collections::hash_map::Entry, rc::Rc};
 
-use crate::{ScopeDbgContext, is_static_value};
+use crate::{ScopeDbgContext, is_static_value, map_rir_literal_to_eval_value};
 
 /// Struct that keeps track of the active RIR blocks (where RIR instructions are added) and the active scopes (which
 /// correspond to the Q#'s program call stack).
@@ -117,6 +117,9 @@ pub struct Scope {
     hybrid_vars: FxHashMap<LocalVarId, Value>,
     /// Maps variable IDs to static literal values, if any.
     static_vars: FxHashMap<VariableId, Literal>,
+    /// Set of mutable local variables in this scope. Used to decide which variables are tracked across
+    /// the partial eval <-> eval boundary (hybrid map <-> classical env).
+    mutable_vars: FxHashSet<LocalVarId>,
     /// Number of currently active blocks (starting from where this scope was created).
     active_block_count: usize,
     /// Debug context, used for generating debug metadata.
@@ -192,6 +195,7 @@ impl Scope {
             active_block_count: 1,
             hybrid_vars,
             static_vars: FxHashMap::default(),
+            mutable_vars: FxHashSet::default(),
             dbg_context: ScopeDbgContext::default(),
             arrays: arrays.into_iter().map(|array| (array, None)).collect(),
         }
@@ -235,6 +239,66 @@ impl Scope {
             .value
     }
 
+    /// Updates the classical local variable values based on the current hybrid local variable values.
+    pub fn update_classical_locals_from_hybrid_locals(&mut self) {
+        for local_var_id in &self.mutable_vars {
+            if self.env.get(*local_var_id).is_none() {
+                continue;
+            }
+            let hybrid_value = self
+                .hybrid_vars
+                .get(local_var_id)
+                .expect("hybrid value should exist");
+            let update_value = match hybrid_value {
+                Value::Var(hybrid_var) => {
+                    // Check to see if there is a static literal value currently tracked for this variable,
+                    // and if so, use it to update the classical local variable value instead.
+                    if let Some(literal) = self.get_static_value(hybrid_var.id.into())
+                        && let Some(literal_val) =
+                            map_rir_literal_to_eval_value(*literal, hybrid_var.ty)
+                    {
+                        literal_val
+                    } else {
+                        hybrid_value.clone()
+                    }
+                }
+                _ => hybrid_value.clone(),
+            };
+            if let Some(classical_var) = self.env.get_mut(*local_var_id) {
+                classical_var.value = update_value;
+            }
+        }
+    }
+
+    /// Collects the local variable values that have been updated based on the current hybrid local variable values,
+    /// allowing callers to choose how to update those in tracked mapping themselves.
+    pub fn collect_updated_local_values(&mut self) -> Vec<(LocalVarId, Value)> {
+        let mut updated_values = Vec::new();
+        for local_var_id in &self.mutable_vars {
+            let hybrid_value = self
+                .hybrid_vars
+                .get(local_var_id)
+                .expect("hybrid value should exist");
+            if let Value::Var(hybrid_var) = hybrid_value
+                && let Some(literal) = self.get_static_value(hybrid_var.id.into())
+                && let Some(literal_val) = map_rir_literal_to_eval_value(*literal, hybrid_var.ty)
+            {
+                // If there is a tracked static literal value for this hybrid variable, use that to check for changes
+                // against the classical local variable value instead of the variable itself.
+                if let Some(classical_value) = self.env.get(*local_var_id).map(|var| &var.value)
+                    && classical_value != &literal_val
+                {
+                    updated_values.push((*local_var_id, classical_value.clone()));
+                }
+            } else if let Some(classical_value) = self.env.get(*local_var_id).map(|var| &var.value)
+                && classical_value != hybrid_value
+            {
+                updated_values.push((*local_var_id, classical_value.clone()));
+            }
+        }
+        updated_values
+    }
+
     /// Gets the value of a hybrid local variable.
     pub fn get_hybrid_local_value(&self, local_var_id: LocalVarId) -> &Value {
         self.hybrid_vars
@@ -252,6 +316,12 @@ impl Scope {
         self.static_vars.insert(var_id, literal);
     }
 
+    // Inserts a local variable into the set of mutable variables so it will be tracked
+    // for updates across evaluation boundaries.
+    pub fn insert_mutable_var(&mut self, local_var_id: LocalVarId) {
+        self.mutable_vars.insert(local_var_id);
+    }
+
     /// Determines whether we are currently evaluating a branch within the scope.
     pub fn is_currently_evaluating_branch(&self) -> bool {
         self.active_block_count > 1

source/compiler/qsc_partial_eval/src/lib.rs

@@ -353,7 +353,14 @@ impl<'a> PartialEvaluator<'a> {
     fn bind_value_to_ident(&mut self, mutability: Mutability, ident: &Ident, value: Value) {
         // We do slightly different things depending on the mutability of the identifier.
         match mutability {
-            Mutability::Mutable => self.bind_value_to_mutable_ident(ident, value),
+            Mutability::Mutable => {
+                // Make sure to track this mutable variable in the current scope to inform
+                // the partial evaluator that it should be monitored for updates across evaluation boundaries.
+                self.eval_context
+                    .get_current_scope_mut()
+                    .insert_mutable_var(ident.id);
+                self.bind_value_to_mutable_ident(ident, value);
+            }
             Mutability::Immutable => {
                 let current_scope = self.eval_context.get_current_scope();
                 if matches!(value, Value::Var(var) if current_scope.get_static_value(var.id.into()).is_none())
@@ -1263,6 +1270,7 @@ impl<'a> PartialEvaluator<'a> {
         let expr = self.package_store.get_expr(store_expr_id);
         let scope_exec_graph = self.get_current_scope_exec_graph().clone();
         let scope = self.eval_context.get_current_scope_mut();
+        scope.update_classical_locals_from_hybrid_locals();
         let exec_graph = scope_exec_graph.get_range(&expr.exec_graph_range);
         let mut state = State::new(
             current_package_id,
@@ -1297,17 +1305,18 @@ impl<'a> PartialEvaluator<'a> {
             Err((error, _)) => Err(Error::from(error)),
         };
 
-        // If this was an assign expression, update the bindings in the hybrid side to keep them in sync and to insert
-        // store instructions for variables of type `Bool`, `Int` or `Double`.
-        if let Ok(EvalControlFlow::Continue(_)) = eval_result {
-            let expr = self.get_expr(expr_id);
-            if let ExprKind::Assign(lhs_expr_id, _)
-            | ExprKind::AssignField(lhs_expr_id, _, _)
-            | ExprKind::AssignIndex(lhs_expr_id, _, _)
-            | ExprKind::AssignOp(_, lhs_expr_id, _) = &expr.kind
-            {
-                self.update_hybrid_bindings_from_classical_bindings(*lhs_expr_id)?;
-            }
+        // Static expressions or sub-expressions may have updated variable values, so we need to collect and apply those updates
+        // to the hybrid maps (potentially updating static variable mappings as well), collecting the pending updates and
+        // making copies of values as needed.
+        let pending_updates = self
+            .eval_context
+            .get_current_scope_mut()
+            .collect_updated_local_values()
+            .into_iter()
+            .map(|(local_var_id, new_value)| (local_var_id, self.copy_value_if_needed(new_value)))
+            .collect::<Vec<_>>();
+        for (local_var_id, new_value) in pending_updates {
+            self.update_hybrid_local(local_var_id, new_value)?;
         }
 
         eval_result
@@ -4150,7 +4159,7 @@ impl<'a> PartialEvaluator<'a> {
                 //
                 // Since expressions call expressions to the `Length` intrinsic will be offloaded to the evaluator,
                 // the evaluator environment also needs to track some non-classical variables.
-                self.update_hybrid_local(lhs_expr, *local_var_id, value.clone())?;
+                self.update_hybrid_local(*local_var_id, value.clone())?;
                 self.update_classical_local(*local_var_id, value);
             }
             (ExprKind::Tuple(exprs), Value::Tuple(values, _)) => {
@@ -4185,12 +4194,7 @@ impl<'a> PartialEvaluator<'a> {
             .update_variable_in_top_frame(local_var_id, value);
     }
 
-    fn update_hybrid_local(
-        &mut self,
-        local_expr: &Expr,
-        local_var_id: LocalVarId,
-        value: Value,
-    ) -> Result<(), Error> {
+    fn update_hybrid_local(&mut self, local_var_id: LocalVarId, value: Value) -> Result<(), Error> {
         let bound_value = self
             .eval_context
             .get_current_scope()
@@ -4215,55 +4219,13 @@ impl<'a> PartialEvaluator<'a> {
                 }
             }
         } else {
-            // Verify that we are not updating a value that does not have a backing variable from a dynamic branch
-            // because it is unsupported.
-            if self
-                .eval_context
-                .get_current_scope()
-                .is_currently_evaluating_branch()
-            {
-                let error_message = format!(
-                    "re-assignment within a dynamic branch is unsupported for type {}",
-                    local_expr.ty
-                );
-                let error =
-                    Error::Unexpected(error_message, self.get_expr_package_span(local_expr.id));
-                return Err(error);
-            }
             self.eval_context
                 .get_current_scope_mut()
                 .update_hybrid_local_value(local_var_id, value);
         }
         Ok(())
     }
 
-    fn update_hybrid_bindings_from_classical_bindings(
-        &mut self,
-        lhs_expr_id: ExprId,
-    ) -> Result<(), Error> {
-        let lhs_expr = &self.get_expr(lhs_expr_id);
-        match &lhs_expr.kind {
-            ExprKind::Hole => {
-                // Nothing to bind to.
-            }
-            ExprKind::Var(Res::Local(local_var_id), _) => {
-                let classical_value = self
-                    .eval_context
-                    .get_current_scope()
-                    .get_classical_local_value(*local_var_id)
-                    .clone();
-                self.update_hybrid_local(lhs_expr, *local_var_id, classical_value)?;
-            }
-            ExprKind::Tuple(exprs) => {
-                for expr_id in exprs {
-                    self.update_hybrid_bindings_from_classical_bindings(*expr_id)?;
-                }
-            }
-            _ => unreachable!("unassignable pattern should be disallowed by compiler"),
-        }
-        Ok(())
-    }
-
     fn value_into_static_mapping(&self, value: Value) -> Value {
         match value {
             Value::Var(var) => {

source/compiler/qsc_partial_eval/src/tests/loops.rs

@@ -305,10 +305,8 @@ fn unitary_call_within_a_repeat_until_loop_unrolled() {
                 Variable(1, Boolean) = Store Bool(true)
                 Call id(2), args( Qubit(0), )
                 Variable(0, Integer) = Store Integer(1)
-                Variable(1, Boolean) = Store Bool(true)
                 Call id(2), args( Qubit(0), )
                 Variable(0, Integer) = Store Integer(2)
-                Variable(1, Boolean) = Store Bool(true)
                 Call id(2), args( Qubit(0), )
                 Variable(0, Integer) = Store Integer(3)
                 Variable(1, Boolean) = Store Bool(false)
@@ -1334,10 +1332,8 @@ fn rotation_call_within_a_repeat_until_loop_unrolled() {
                 Variable(1, Boolean) = Store Bool(true)
                 Call id(2), args( Double(0), Qubit(0), )
                 Variable(0, Integer) = Store Integer(1)
-                Variable(1, Boolean) = Store Bool(true)
                 Call id(2), args( Double(1), Qubit(0), )
                 Variable(0, Integer) = Store Integer(2)
-                Variable(1, Boolean) = Store Bool(true)
                 Call id(2), args( Double(2), Qubit(0), )
                 Variable(0, Integer) = Store Integer(3)
                 Variable(1, Boolean) = Store Bool(false)
@@ -2049,6 +2045,7 @@ fn classical_while_inside_dynamic_while_folds_mutable_variable() {
                 Return Integer(0)
             Block 3:Block:
                 Variable(3, Integer) = Store Integer(0)
+                Variable(3, Integer) = Store Integer(3)
                 Variable(4, Integer) = Store Variable(0, Integer)
                 Variable(5, Integer) = Add Variable(4, Integer), Integer(3)
                 Variable(0, Integer) = Store Variable(5, Integer)

source/compiler/qsc_partial_eval/src/tests/misc.rs

@@ -169,17 +169,12 @@ fn unitary_call_within_an_if_with_classical_condition_within_a_repeat_until_loop
                 Variable(1, Boolean) = Store Bool(true)
                 Call id(2), args( Qubit(0), )
                 Variable(0, Integer) = Store Integer(1)
-                Variable(1, Boolean) = Store Bool(true)
                 Variable(0, Integer) = Store Integer(2)
-                Variable(1, Boolean) = Store Bool(true)
                 Call id(2), args( Qubit(0), )
                 Variable(0, Integer) = Store Integer(3)
-                Variable(1, Boolean) = Store Bool(true)
                 Variable(0, Integer) = Store Integer(4)
-                Variable(1, Boolean) = Store Bool(true)
                 Call id(2), args( Qubit(0), )
                 Variable(0, Integer) = Store Integer(5)
-                Variable(1, Boolean) = Store Bool(true)
                 Variable(0, Integer) = Store Integer(6)
                 Variable(1, Boolean) = Store Bool(false)
                 Call id(3), args( Integer(0), Tag(0, 3), )
@@ -1135,3 +1130,50 @@ fn test_dynamic_expr_including_bigint_binop() {
                 Return Integer(0)"#]],
     );
 }
+
+#[test]
+fn tuple_swap_update_of_dynamic_mutable_variables_in_static_subexpr() {
+    let program = get_rir_program(indoc! {
+        r#"
+        operation Main() : (Int, Int) {
+            use q = Qubit();
+            X(q);
+            let d = MResetZ(q) == One;
+            mutable x = 0;
+            mutable y = 0;
+            let s = d and { set x = 3; true };
+            let _ = { set (x, y) = (y, x); true };
+            (x, y)
+        }
+        "#,
+    });
+    assert_blocks(
+        &program,
+        &expect![[r#"
+            Blocks:
+            Block 0:Block:
+                Call id(1), args( Pointer, )
+                Call id(2), args( Qubit(0), )
+                Call id(3), args( Qubit(0), Result(0), )
+                Variable(0, Boolean) = Call id(4), args( Result(0), )
+                Variable(1, Boolean) = Store Variable(0, Boolean)
+                Variable(2, Boolean) = Store Variable(1, Boolean)
+                Variable(3, Integer) = Store Integer(0)
+                Variable(4, Integer) = Store Integer(0)
+                Variable(5, Boolean) = Store Variable(2, Boolean)
+                Variable(6, Boolean) = Store Bool(false)
+                Branch Variable(5, Boolean), 2, 1
+            Block 1:Block:
+                Variable(7, Boolean) = Store Variable(6, Boolean)
+                Variable(4, Integer) = Store Integer(3)
+                Variable(3, Integer) = Store Integer(0)
+                Call id(5), args( Integer(2), Tag(0, 3), )
+                Call id(6), args( Integer(0), Tag(1, 5), )
+                Call id(6), args( Integer(3), Tag(2, 5), )
+                Return Integer(0)
+            Block 2:Block:
+                Variable(3, Integer) = Store Integer(3)
+                Variable(6, Boolean) = Store Bool(true)
+                Jump(1)"#]],
+    );
+}

source/qdk_package/tests-integration/resources/adaptive/input/MutationAcrossClassicalBranches.qs

@@ -0,0 +1,13 @@
+namespace Test {
+    operation Main() : Int {
+        use q = Qubit();
+        mutable n = 0;
+        if true {
+            set n = 5;
+        }
+        if MResetZ(q) == One {
+            set n = 1;
+        }
+        n
+    }
+}