From cd982d8d8e05d2740f9aec87a22b9b5fd9c47078 Mon Sep 17 00:00:00 2001
From: Mathieu Poumeyrol <kali@zoy.org>
Date: Mon, 9 Dec 2024 11:02:51 +0100
Subject: [PATCH] exp tdim caching

---
 Cargo.toml          |  1 +
 data/Cargo.toml     |  3 ++-
 data/src/dim/sym.rs | 62 ++++++++++++++++++++++++---------------------
 3 files changed, 36 insertions(+), 30 deletions(-)

diff --git a/Cargo.toml b/Cargo.toml
index b66d7604fa..bf8b07a7a9 100644
--- a/Cargo.toml
+++ b/Cargo.toml
@@ -172,6 +172,7 @@ prost-types = "0.11.0"
 py_literal = "0.4.0"
 rand = { version = "0.8.4", features = ["small_rng"] }
 rand_distr = "0.4"
+rapidhash = "1.2"
 rayon = "1.10"
 readings-probe = "0.1.3"
 regex = "1.5.4"
diff --git a/data/Cargo.toml b/data/Cargo.toml
index 69a225787d..00c4bfcc05 100644
--- a/data/Cargo.toml
+++ b/data/Cargo.toml
@@ -9,7 +9,7 @@ keywords = [ "TensorFlow", "NeuralNetworks" ]
 categories = [ "science" ]
 autobenches = false
 edition = "2021"
-rust-version = "1.75"
+rust-version = "1.77"
 
 [badges]
 maintenance = { status = "actively-developed" }
@@ -27,6 +27,7 @@ nom.workspace = true
 num-complex = { workspace = true, optional = true }
 num-integer.workspace = true
 num-traits.workspace = true
+rapidhash.workspace = true
 smallvec.workspace = true
 lazy_static.workspace = true
 scan_fmt.workspace = true
diff --git a/data/src/dim/sym.rs b/data/src/dim/sym.rs
index 77415187df..ad48bd6569 100644
--- a/data/src/dim/sym.rs
+++ b/data/src/dim/sym.rs
@@ -28,6 +28,7 @@ pub struct SymbolScopeData {
     table: DefaultStringInterner,
     assertions: Vec<Assertion>,
     scenarios: Vec<(String, Vec<Assertion>)>,
+    cache_proven_positive_or_zero: RefCell<rapidhash::RapidHashMap<TDim, bool>>,
 }
 
 impl SymbolScope {
@@ -147,7 +148,7 @@ impl SymbolScope {
         let locked = self.0.lock();
         let locked = locked.borrow();
         if locked.scenarios.len() == 0 {
-            return Ok(None)
+            return Ok(None);
         }
         let mut maybe = None;
         for (ix, (_name, assertions)) in locked.scenarios.iter().enumerate() {
@@ -158,7 +159,7 @@ impl SymbolScope {
             } else if maybe.is_none() {
                 maybe = Some(ix);
             } else {
-                return Ok(None)
+                return Ok(None);
             }
         }
         if maybe.is_some() {
@@ -201,38 +202,41 @@ impl SymbolScopeData {
         if let TDim::Val(v) = t {
             return *v >= 0;
         }
-        let positives = self.assertions.iter().filter_map(|i| i.as_known_positive()).collect_vec();
-        let mut visited = vec![];
-        let mut todo = vec![t.clone()];
-        while let Some(t) = todo.pop() {
-            if t.to_i64().is_ok_and(|i| i >= 0) {
-                return true;
-            }
-            if t.inclusive_bound(self, false).is_some_and(|l| l >= 0) {
-                return true;
-            }
-            let syms = t.symbols();
-            for s in syms {
-                let me = t.guess_slope(&s);
-                for pos in &positives {
-                    if pos.symbols().contains(&s) {
-                        let other = pos.guess_slope(&s);
-                        if me.0.signum() == other.0.signum() {
-                            let new = t.clone() * me.1 * other.0.abs()
-                                - pos.clone() * me.0.abs() * other.1;
-                            if !visited.contains(&new) {
-                                todo.push(new);
+        *self.cache_proven_positive_or_zero.borrow_mut().entry(t.clone()).or_insert_with(|| {
+            let positives =
+                self.assertions.iter().filter_map(|i| i.as_known_positive()).collect_vec();
+            let mut visited = vec![];
+            let mut todo = vec![t.clone()];
+            while let Some(t) = todo.pop() {
+                if t.to_i64().is_ok_and(|i| i >= 0) {
+                    return true;
+                }
+                if t.inclusive_bound(self, false).is_some_and(|l| l >= 0) {
+                    return true;
+                }
+                let syms = t.symbols();
+                for s in syms {
+                    let me = t.guess_slope(&s);
+                    for pos in &positives {
+                        if pos.symbols().contains(&s) {
+                            let other = pos.guess_slope(&s);
+                            if me.0.signum() == other.0.signum() {
+                                let new = t.clone() * me.1 * other.0.abs()
+                                    - pos.clone() * me.0.abs() * other.1;
+                                if !visited.contains(&new) {
+                                    todo.push(new);
+                                }
                             }
                         }
                     }
                 }
+                visited.push(t);
+                if visited.len() > 10 {
+                    break;
+                }
             }
-            visited.push(t);
-            if visited.len() > 10 {
-                break;
-            }
-        }
-        false
+            false
+        })
     }
 }