From 5a83316701f2b79f0ada3b76a85463f14935aa52 Mon Sep 17 00:00:00 2001
From: Michal Moskal <michal@moskal.me>
Date: Thu, 14 Mar 2024 21:18:12 +0000
Subject: [PATCH] fix grammar opt

---
 controllers/aici_abi/src/earley/grammar.rs | 31 +++++++++++++---------
 1 file changed, 18 insertions(+), 13 deletions(-)

diff --git a/controllers/aici_abi/src/earley/grammar.rs b/controllers/aici_abi/src/earley/grammar.rs
index 8a3543ea..2143cd56 100644
--- a/controllers/aici_abi/src/earley/grammar.rs
+++ b/controllers/aici_abi/src/earley/grammar.rs
@@ -126,6 +126,21 @@ impl Grammar {
         }
     }
 
+    fn rule_shape(&self, r: &Rule) -> Vec<Option<SymIdx>> {
+        let mut shape = Vec::new();
+        let mut had_term = false;
+        for s in &r.rhs {
+            let sym = self.sym_data(*s);
+            if !had_term && sym.is_terminal() {
+                had_term = true;
+                shape.push(None);
+            } else {
+                shape.push(Some(*s));
+            }
+        }
+        shape
+    }
+
     fn collapse_terminals(&self) -> Self {
         let mut outp = Grammar::new();
         for sym in &self.symbols {
@@ -134,19 +149,8 @@ impl Grammar {
             }
             let mut rules_by_shape = FxHashMap::default();
             for rule in &sym.rules {
-                let shape = rule
-                    .rhs
-                    .iter()
-                    .map(|s| {
-                        if self.sym_data(*s).is_terminal() {
-                            None
-                        } else {
-                            Some(*s)
-                        }
-                    })
-                    .collect::<Vec<_>>();
                 rules_by_shape
-                    .entry(shape)
+                    .entry(self.rule_shape(rule))
                     .or_insert_with(Vec::new)
                     .push(rule);
             }
@@ -280,6 +284,7 @@ impl Grammar {
 
 impl Debug for Grammar {
     fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+        writeln!(f, "Grammar:")?;
         for sym in &self.symbols {
             match sym.bytes {
                 Some(ref bytes) if sym.name.starts_with("T@") => {
@@ -589,5 +594,5 @@ fn rule_to_string(lhs: &str, mut rhs: Vec<SymName>, dot: Option<usize>) -> Strin
             }
         }
     }
-    format!("{} ::= {}", lhs, outp.join(" "))
+    format!("{:15} ⇦ {}", lhs, outp.join(" "))
 }