From c7e928499799198f5ec6327c08fdd949470d064a Mon Sep 17 00:00:00 2001
From: Max Willsey <me@mwillsey.com>
Date: Fri, 22 Mar 2024 15:40:39 -0700
Subject: [PATCH] Add way to create numeric variables without parsing

Closes #287.
---
 src/subst.rs | 75 ++++++++++++++++++++++++++++++++++++++++++++++------
 1 file changed, 67 insertions(+), 8 deletions(-)

diff --git a/src/subst.rs b/src/subst.rs
index ffd57cb5..e064d4d4 100644
--- a/src/subst.rs
+++ b/src/subst.rs
@@ -12,12 +12,45 @@ use thiserror::Error;
 ///
 /// [`FromStr`]: std::str::FromStr
 #[derive(Copy, Clone, PartialEq, Eq, PartialOrd, Ord, Hash)]
-pub struct Var(Symbol);
+pub struct Var(VarInner);
+
+impl Var {
+    /// Create a new variable from a u32.
+    ///
+    /// You can also use special syntax `?#3`, `?#42` to denote a numeric variable.
+    /// These avoid some symbol interning, and can also be created manually from
+    /// using this function or the `From` impl.
+    ///
+    /// ```rust
+    /// # use egg::*;
+    /// assert_eq!(Var::from(12), "?#12".parse().unwrap());
+    /// assert_eq!(Var::from_u32(12), "?#12".parse().unwrap());
+    /// ```
+    pub fn from_u32(num: u32) -> Self {
+        Var(VarInner::Num(num))
+    }
+
+    /// If this variable was created from a u32, get it back out.
+    pub fn as_u32(&self) -> Option<u32> {
+        match self.0 {
+            VarInner::Num(num) => Some(num),
+            _ => None,
+        }
+    }
+}
+
+#[derive(Copy, Clone, PartialEq, Eq, PartialOrd, Ord, Hash)]
+enum VarInner {
+    Sym(Symbol),
+    Num(u32),
+}
 
 #[derive(Debug, Error)]
 pub enum VarParseError {
     #[error("pattern variable {0:?} should have a leading question mark")]
     MissingQuestionMark(String),
+    #[error("number pattern variable {0:?} was malformed")]
+    BadNumber(String),
 }
 
 impl FromStr for Var {
@@ -26,23 +59,38 @@ impl FromStr for Var {
     fn from_str(s: &str) -> Result<Self, Self::Err> {
         use VarParseError::*;
 
-        if s.starts_with('?') && s.len() > 1 {
-            Ok(Var(s.into()))
-        } else {
-            Err(MissingQuestionMark(s.to_owned()))
+        match s.as_bytes() {
+            [b'?', b'#', ..] => s[2..]
+                .parse()
+                .map(|num| Var(VarInner::Num(num)))
+                .map_err(|_| BadNumber(s.to_owned())),
+            [b'?', ..] if s.len() > 1 => Ok(Var(VarInner::Sym(Symbol::from(s)))),
+            _ => Err(MissingQuestionMark(s.to_owned())),
         }
     }
 }
 
 impl Display for Var {
     fn fmt(&self, f: &mut Formatter<'_>) -> fmt::Result {
-        Display::fmt(&self.0, f)
+        match self.0 {
+            VarInner::Sym(sym) => write!(f, "{}", sym),
+            VarInner::Num(num) => write!(f, "?#{}", num),
+        }
     }
 }
 
 impl Debug for Var {
     fn fmt(&self, f: &mut Formatter<'_>) -> fmt::Result {
-        Debug::fmt(&self.0, f)
+        match self.0 {
+            VarInner::Sym(sym) => write!(f, "{:?}", sym),
+            VarInner::Num(num) => write!(f, "?#{}", num),
+        }
+    }
+}
+
+impl From<u32> for Var {
+    fn from(num: u32) -> Self {
+        Var(VarInner::Num(num))
     }
 }
 
@@ -87,7 +135,7 @@ impl std::ops::Index<Var> for Subst {
     fn index(&self, var: Var) -> &Self::Output {
         match self.get(var) {
             Some(id) => id,
-            None => panic!("Var '{}={}' not found in {:?}", var.0, var, self),
+            None => panic!("Var '{}={}' not found in {:?}", var, var, self),
         }
     }
 }
@@ -118,5 +166,16 @@ mod tests {
         assert!(Var::from_str("a").is_err());
         assert!(Var::from_str("a?").is_err());
         assert!(Var::from_str("?").is_err());
+        assert!(Var::from_str("?#").is_err());
+        assert!(Var::from_str("?#foo").is_err());
+
+        // numeric vars
+        assert_eq!(Var::from_str("?#0").unwrap(), Var(VarInner::Num(0)));
+        assert_eq!(Var::from_str("?#010").unwrap(), Var(VarInner::Num(10)));
+        assert_eq!(
+            Var::from_str("?#10").unwrap(),
+            Var::from_str("?#0010").unwrap()
+        );
+        assert_eq!(Var::from_str("?#010").unwrap(), Var(VarInner::Num(10)));
     }
 }