more work on scanarios

sonos · Sep 18, 2024 · 9d06ad1 · 9d06ad1
1 parent 370f175
commit 9d06ad1
Show file tree

Hide file tree

Showing 4 changed files with 18 additions and 4 deletions.
diff --git a/data/src/dim/parse.rs b/data/src/dim/parse.rs
@@ -22,7 +22,7 @@ pub fn parse_assertion(symbol_table: &SymbolScope, input: &str) -> TractResult<A
 }
 
 fn assertion<'i>(s: &SymbolScope, i: &'i str) -> IResult<&'i str, Assertion> {
-    alt((
+    delimited(spaces, alt((
         map(separated_pair(|i| expr(s, i), stag("=="), |i| expr(s, i)), |(a, b)| {
             Assertion::Eq(a, b)
         }),
@@ -38,7 +38,7 @@ fn assertion<'i>(s: &SymbolScope, i: &'i str) -> IResult<&'i str, Assertion> {
         map(separated_pair(|i| expr(s, i), stag(">"), |i| expr(s, i)), |(a, b)| {
             Assertion::GT(a, b)
         }),
-    ))(i)
+    )), spaces)(i)
 }
 
 fn expr<'i>(symbol_table: &SymbolScope, i: &'i str) -> IResult<&'i str, TDim> {

diff --git a/data/src/dim/sym.rs b/data/src/dim/sym.rs
@@ -148,6 +148,10 @@ impl SymbolScopeData {
         self.scenarios.keys().map(|s| s.as_ref())
     }
 
+    pub fn scenario(&self, s: &str) -> impl Iterator<Item = &'_ Assertion> {
+        self.scenarios[s].iter()
+    }
+
     pub fn resolving<R>(&self, sym: &Symbol, f: impl FnOnce(&str) -> R) -> Option<R> {
         self.table.resolve(sym.1).map(f)
     }

diff --git a/nnef/src/deser.rs b/nnef/src/deser.rs
@@ -62,7 +62,11 @@ impl<'mb> ModelBuilder<'mb> {
                     self.symbols.push(symbol);
                 }
                 "tract_assert" => {
-                    self.model.symbols.add_assertion(&ext.1)?;
+                    if let Some((scen, rule)) = ext.1.split_once(':') {
+                        self.model.symbols.add_scenario_assertion(scen, rule)?;
+                    } else {
+                        self.model.symbols.add_assertion(&ext.1)?;
+                    }
                 }
                 "KHR_enable_fragment_definitions" | "KHR_enable_operator_expressions" => (),
                 _ => {

diff --git a/nnef/src/ser.rs b/nnef/src/ser.rs
@@ -204,9 +204,15 @@ impl<'a> IntoAst<'a> {
         for sym in self.model.symbols.all_symbols() {
             extension.push(("tract_symbol".into(), sym.to_string()));
         }
-        for assert in self.model.symbols.all_assertions() {
+        let locked = self.model.symbols.0.lock();
+        for assert in locked.borrow().all_assertions() {
             extension.push(("tract_assert".into(), assert.to_string()));
         }
+        for scenario in locked.borrow().scenarios() {
+            for assert in locked.borrow().scenario(scenario) {
+                extension.push(("tract_assert".into(), format!("{scenario}: {assert}")));
+            }
+        }
         let properties = FragmentDef {
             decl: FragmentDecl {
                 id: Identifier("tract_core_properties".to_string()),