From b3468842690fec6e086bbf4ae65ac53663b680f5 Mon Sep 17 00:00:00 2001
From: Mario Carneiro <di.gama@gmail.com>
Date: Sat, 15 Jun 2024 05:14:45 +0200
Subject: [PATCH] support non-compressed proofs in axiom_use.rs

fixes #160
---
 metamath-rs/src/axiom_use.rs | 79 +++++++++++++++++++++++-------------
 1 file changed, 51 insertions(+), 28 deletions(-)

diff --git a/metamath-rs/src/axiom_use.rs b/metamath-rs/src/axiom_use.rs
index a8bf1e7..afe4445 100644
--- a/metamath-rs/src/axiom_use.rs
+++ b/metamath-rs/src/axiom_use.rs
@@ -9,8 +9,8 @@ use crate::segment::SegmentRef;
 use crate::segment_set::SegmentSet;
 use crate::statement::{CommandToken, StatementAddress, TokenPtr};
 use crate::util::HashMap;
-use crate::StatementType;
 use crate::{as_str, database::time, Database};
+use crate::{StatementRef, StatementType};
 
 /// Diagnostics issued when checking axiom usage in the Database.
 ///
@@ -89,21 +89,7 @@ impl<'a> UsagePass<'a> {
                     }
                     StatementType::Provable => {
                         let label = stmt.label();
-                        let mut usage = Bitset::new();
-                        let mut i = 1;
-                        loop {
-                            let tk = stmt.proof_slice_at(i);
-                            i += 1;
-                            if tk == b")" {
-                                break;
-                            }
-                            if let Some(usage2) = self.axiom_use_map.get(tk) {
-                                usage |= usage2
-                            }
-                            if i == stmt.proof_len() {
-                                break;
-                            }
-                        }
+                        let usage = get_proof_usage(&stmt, &self.axiom_use_map);
                         self.axiom_use_map.insert(label, usage);
                     }
                     StatementType::Axiom => {
@@ -144,23 +130,13 @@ impl Database {
             let mut stmt_list = vec![];
             for sref in self.statements().filter(|stmt| stmt.is_assertion()) {
                 let label = sref.label();
-                let mut usage = Bitset::new();
                 if label_test(label) {
+                    let mut usage = Bitset::new();
                     usage.set_bit(stmt_list.len());
                     stmt_list.push(as_str(label));
                     stmt_use_map.insert(label, usage);
                 } else if sref.statement_type() == StatementType::Provable {
-                    let mut i = 1;
-                    loop {
-                        let tk = sref.proof_slice_at(i);
-                        i += 1;
-                        if tk == b")" {
-                            break;
-                        }
-                        if let Some(usage2) = stmt_use_map.get(tk) {
-                            usage |= usage2
-                        }
-                    }
+                    let usage = get_proof_usage(&sref, &stmt_use_map);
                     write!(out, "{}:", as_str(label))?;
                     for i in &usage {
                         write!(out, " {}", stmt_list[i])?;
@@ -173,3 +149,50 @@ impl Database {
         })
     }
 }
+
+#[inline]
+fn get_proof_usage(sref: &StatementRef<'_>, stmt_use_map: &HashMap<&[u8], Bitset>) -> Bitset {
+    let mut usage = Bitset::new();
+    if sref.proof_slice_at(0) == b"(" {
+        let mut i = 1;
+        loop {
+            let tk = sref.proof_slice_at(i);
+            i += 1;
+            if tk == b")" {
+                break;
+            }
+            if let Some(usage2) = stmt_use_map.get(tk) {
+                usage |= usage2
+            }
+            if i == sref.proof_len() {
+                break;
+            }
+        }
+    } else {
+        #[cold]
+        fn process_non_compressed(
+            sref: &StatementRef<'_>,
+            stmt_use_map: &HashMap<&[u8], Bitset>,
+            usage: &mut Bitset,
+        ) {
+            let mut i = 0;
+            while i < sref.proof_len() {
+                let tk = sref.proof_slice_at(i);
+                if let Some(usage2) = stmt_use_map.get(tk) {
+                    *usage |= usage2
+                } else if let Some(i) = tk
+                    .iter()
+                    .position(|&x| x == b'=')
+                    .or_else(|| tk.iter().position(|&x| x == b':'))
+                {
+                    if let Some(usage2) = stmt_use_map.get(&tk[i + 1..]) {
+                        *usage |= usage2
+                    }
+                }
+                i += 1;
+            }
+        }
+        process_non_compressed(sref, stmt_use_map, &mut usage)
+    }
+    usage
+}