From 3520c49aa3d71a0038e3bc2262bef0ea2ea51ad9 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Kevin=20L=C3=A4ufer?= <laeufer@cornell.edu>
Date: Mon, 25 Nov 2024 09:46:09 -0500
Subject: [PATCH] ts: make get_all_expressions a method

---
 src/system/transform.rs         | 24 +-----------------------
 src/system/transition_system.rs | 25 ++++++++++++++++++++++++-
 2 files changed, 25 insertions(+), 24 deletions(-)

diff --git a/src/system/transform.rs b/src/system/transform.rs
index 35fd303..6a03eb1 100644
--- a/src/system/transform.rs
+++ b/src/system/transform.rs
@@ -48,7 +48,7 @@ pub fn do_transform(
     mode: ExprTransformMode,
     tran: impl FnMut(&mut Context, ExprRef, &[ExprRef]) -> Option<ExprRef>,
 ) {
-    let todo = get_root_expressions(sys);
+    let todo = sys.get_all_expressions();
     let mut transformed = DenseExprMetaData::default();
     do_transform_expr(ctx, mode, &mut transformed, todo, tran);
 
@@ -61,25 +61,3 @@ pub fn do_transform(
         }
     });
 }
-
-fn get_root_expressions(sys: &TransitionSystem) -> Vec<ExprRef> {
-    // include all input, output, assertion and assumptions expressions
-    let mut out = vec![];
-    out.extend_from_slice(sys.inputs.as_slice());
-    out.extend(sys.outputs.iter().map(|o| o.expr));
-    out.extend_from_slice(sys.bad_states.as_slice());
-    out.extend_from_slice(sys.constraints.as_slice());
-
-    // include all states
-    for state in sys.states.iter() {
-        out.push(state.symbol);
-        if let Some(init) = state.init {
-            out.push(init);
-        }
-        if let Some(next) = state.next {
-            out.push(next);
-        }
-    }
-
-    out
-}
diff --git a/src/system/transition_system.rs b/src/system/transition_system.rs
index 932d581..b9f58ed 100644
--- a/src/system/transition_system.rs
+++ b/src/system/transition_system.rs
@@ -110,7 +110,7 @@ impl TransitionSystem {
         HashMap::from_iter(self.states.iter().map(|s| (s.symbol, s)))
     }
 
-    /// Update all output, input, assume, assert, `state.next` and `state.init` expressions.
+    /// Update all output, input, assume, assert, state expressions.
     /// If `update` returns `None`, no update is performed.
     pub fn update_expressions(&mut self, mut update: impl FnMut(ExprRef) -> Option<ExprRef>) {
         for old in self.inputs.iter_mut() {
@@ -131,4 +131,27 @@ impl TransitionSystem {
             state.next = state.next.and_then(|old| update(old));
         }
     }
+
+    /// Returns a list of all output, input, assume, assert and state expressions.
+    pub fn get_all_expressions(&self) -> Vec<ExprRef> {
+        // include all input, output, assertion and assumptions expressions
+        let mut out = vec![];
+        out.extend_from_slice(self.inputs.as_slice());
+        out.extend(self.outputs.iter().map(|o| o.expr));
+        out.extend_from_slice(self.bad_states.as_slice());
+        out.extend_from_slice(self.constraints.as_slice());
+
+        // include all states
+        for state in self.states.iter() {
+            out.push(state.symbol);
+            if let Some(init) = state.init {
+                out.push(init);
+            }
+            if let Some(next) = state.next {
+                out.push(next);
+            }
+        }
+
+        out
+    }
 }