From 76940d36aca537fc618bd7c255a0e32193afaa4b Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Kevin=20L=C3=A4ufer?= <laeufer@cs.berkeley.edu>
Date: Thu, 14 Dec 2023 12:21:05 -0500
Subject: [PATCH] synth: correctly detect synthesis variables in submodules

---
 synth/src/repair.rs | 44 ++++++++++++++++++++++++++++++++++----------
 1 file changed, 34 insertions(+), 10 deletions(-)

diff --git a/synth/src/repair.rs b/synth/src/repair.rs
index 0130dbb..64ddc08 100644
--- a/synth/src/repair.rs
+++ b/synth/src/repair.rs
@@ -137,16 +137,19 @@ impl RepairVars {
         let mut free = Vec::new();
 
         for state in sys.states() {
-            let name = state.symbol.get_symbol_name(ctx).unwrap();
-            if name.starts_with(SYNTH_CHANGE_PREFIX) {
-                assert_eq!(
-                    state.symbol.get_bv_type(ctx).unwrap(),
-                    1,
-                    "all change variables need to be boolean"
-                );
-                change.push(state.symbol);
-            } else if name.starts_with(SYNTH_VAR_PREFIX) {
-                free.push(state.symbol);
+            match classify_state(ctx, state) {
+                StateType::ChangeVar => {
+                    assert_eq!(
+                        state.symbol.get_bv_type(ctx).unwrap(),
+                        1,
+                        "all change variables need to be boolean"
+                    );
+                    change.push(state.symbol);
+                }
+                StateType::FreeVar => {
+                    free.push(state.symbol);
+                }
+                StateType::Other => {} // nothing to do
             }
         }
 
@@ -267,6 +270,27 @@ impl RepairVars {
     }
 }
 
+enum StateType {
+    ChangeVar,
+    FreeVar,
+    Other,
+}
+
+/// Determines whether a state is a synthesis variable and what kind by looking at the name.
+fn classify_state(ctx: &Context, state: &State) -> StateType {
+    let name = state.symbol.get_symbol_name(ctx).unwrap();
+    let suffix = name.split('.').last().unwrap();
+    // important to check the change prefix first
+    // (since the var prefix is a prefix of the change prefix)
+    if suffix.starts_with(SYNTH_CHANGE_PREFIX) {
+        StateType::ChangeVar
+    } else if suffix.starts_with(SYNTH_VAR_PREFIX) {
+        StateType::FreeVar
+    } else {
+        StateType::Other
+    }
+}
+
 #[derive(Debug, Clone)]
 pub struct RepairAssignment {
     pub change: Vec<bool>,