feat: better interfaces

ModelTC · Oct 26, 2023 · 2748dc6 · 2748dc6
1 parent 5e19bb5
commit 2748dc6
Show file tree

Hide file tree

Showing 3 changed files with 112 additions and 125 deletions.
diff --git a/src/sam/mod.rs b/src/sam/mod.rs
@@ -1,12 +1,12 @@
 mod state;
 pub use state::GeneralSAMState;
 
-use std::{
-    collections::{BTreeMap, VecDeque},
-    convert::Infallible,
-};
+use std::{collections::BTreeMap, convert::Infallible};
 
-use crate::trie_alike::{IterAsChain, TravelEvent, TrieNodeAlike};
+use crate::{
+    trie_alike::{IterAsChain, TrieNodeAlike},
+    TravelEvent,
+};
 
 pub type GeneralSAMNodeID = usize;
 pub const SAM_NIL_NODE_ID: GeneralSAMNodeID = 0;
@@ -142,22 +142,14 @@ impl<T: Ord + Clone> GeneralSAM<T> {
     where
         TN::InnerType: Into<T>,
     {
-        let mut queue = VecDeque::new();
-        let mut last_node_id = SAM_ROOT_NODE_ID;
-        node.bfs_travel(|event| -> Result<(), Infallible> {
+        node.bfs_travel(|event| -> Result<GeneralSAMNodeID, Infallible> {
             match event {
-                TravelEvent::Push(_, None) => {
-                    queue.push_back(SAM_ROOT_NODE_ID);
-                }
-                TravelEvent::Pop(_) => {
-                    last_node_id = queue.pop_front().unwrap();
+                TravelEvent::PushRoot(_) => Ok(SAM_ROOT_NODE_ID),
+                TravelEvent::Push(cur_tn, cur_node_id, key) => {
+                    Ok(self.insert_node_trans(*cur_node_id, key, cur_tn.is_accepting()))
                 }
-                TravelEvent::Push(tn, Some(key)) => {
-                    let new_node_id = self.insert_node_trans(last_node_id, key, tn.is_accepting());
-                    queue.push_back(new_node_id);
-                }
-            };
-            Ok(())
+                TravelEvent::Pop(_, cur_node_id) => Ok(cur_node_id),
+            }
         })
         .unwrap();
     }

diff --git a/src/sam/state.rs b/src/sam/state.rs
@@ -1,5 +1,3 @@
-use std::collections::VecDeque;
-
 use crate::trie_alike::{TravelEvent, TrieNodeAlike};
 
 use super::{GeneralSAM, GeneralSAMNode, SAM_NIL_NODE_ID, SAM_ROOT_NODE_ID};
@@ -16,13 +14,13 @@ impl<'s> GeneralSAMState<'s, u8> {
     }
 }
 
-impl<'s> GeneralSAMState<'s, char> {
-    pub fn feed_chars(self, seq: &'s str) -> Self {
+impl GeneralSAMState<'_, char> {
+    pub fn feed_chars(self, seq: &str) -> Self {
         self.feed(seq.chars())
     }
 }
 
-impl<'s, T: Ord + Clone> GeneralSAMState<'s, T> {
+impl<T: Ord + Clone> GeneralSAMState<'_, T> {
     pub fn is_nil(&self) -> bool {
         self.node_id == SAM_NIL_NODE_ID
     }
@@ -58,116 +56,98 @@ impl<'s, T: Ord + Clone> GeneralSAMState<'s, T> {
             }
     }
 
-    pub fn feed_ref<Seq: IntoIterator<Item = &'s T>>(self, seq: Seq) -> Self {
-        self.feed_ref_iter(seq.into_iter())
+    pub fn feed<Seq: IntoIterator<Item = T>>(self, seq: Seq) -> Self {
+        self.feed_iter(seq.into_iter())
     }
 
-    pub fn feed_ref_iter<Iter: Iterator<Item = &'s T>>(mut self, iter: Iter) -> Self {
+    pub fn feed_iter<Iter: Iterator<Item = T>>(mut self, iter: Iter) -> Self {
         for t in iter {
             if self.is_nil() {
                 break;
             }
-            self.goto(t)
+            self.goto(&t)
         }
         self
     }
+}
 
-    pub fn feed<Seq: IntoIterator<Item = T>>(self, seq: Seq) -> Self {
-        self.feed_iter(seq.into_iter())
+impl<'s, T: Ord + Clone> GeneralSAMState<'s, T> {
+    pub fn feed_ref<Seq: IntoIterator<Item = &'s T>>(self, seq: Seq) -> Self {
+        self.feed_ref_iter(seq.into_iter())
     }
 
-    pub fn feed_iter<Iter: Iterator<Item = T>>(mut self, iter: Iter) -> Self {
+    pub fn feed_ref_iter<Iter: Iterator<Item = &'s T>>(mut self, iter: Iter) -> Self {
         for t in iter {
             if self.is_nil() {
                 break;
             }
-            self.goto(&t)
+            self.goto(t)
         }
         self
     }
+}
 
-    pub fn bfs_along<
-        TN: TrieNodeAlike<InnerType = T> + Sized,
-        E,
-        F: FnMut(TravelEvent<(GeneralSAMState<'_, T>, &TN), TN::InnerType>) -> Result<(), E>,
+impl<'s, T: Ord + Clone> GeneralSAMState<'s, T> {
+    fn wrap_travel_along_callback<
+        TN: TrieNodeAlike<InnerType = T>,
+        ExtraType,
+        ErrorType,
+        F: 's
+            + FnMut(
+                TravelEvent<(&GeneralSAMState<T>, &TN), ExtraType, TN::InnerType>,
+            ) -> Result<ExtraType, ErrorType>,
     >(
-        &self,
-        trie_node: TN,
+        &'s self,
         mut callback: F,
-    ) -> Result<(), E> {
-        let mut queue = VecDeque::new();
-        let mut cur_node_id = self.node_id;
-
-        trie_node.bfs_travel(|event| match event {
-            TravelEvent::Push(tn, Some(key)) => {
-                let next_node_id = self
-                    .sam
-                    .node_pool
-                    .get(cur_node_id)
-                    .and_then(|x| x.trans.get(&key).copied())
-                    .unwrap_or(SAM_NIL_NODE_ID);
-                callback(TravelEvent::Push(
-                    (self.sam.get_state(next_node_id), tn),
-                    Some(key),
-                ))?;
-                queue.push_back(next_node_id);
-                Ok(())
+    ) -> impl FnMut(
+        TravelEvent<&TN, (GeneralSAMState<'s, T>, ExtraType), TN::InnerType>,
+    ) -> Result<(GeneralSAMState<'s, T>, ExtraType), ErrorType> {
+        move |event| match event {
+            TravelEvent::PushRoot(trie_root) => {
+                let res = callback(TravelEvent::PushRoot((self, trie_root)))?;
+                Ok((self.clone(), res))
             }
-            TravelEvent::Push(tn, None) => {
-                callback(TravelEvent::Push(
-                    (self.sam.get_state(self.node_id), tn),
-                    None,
-                ))?;
-                queue.push_back(self.node_id);
-                Ok(())
+            TravelEvent::Push(cur_tn, (cur_state, cur_extra), key) => {
+                let mut next_state = cur_state.clone();
+                next_state.goto(&key);
+                let next_extra =
+                    callback(TravelEvent::Push((&next_state, &cur_tn), cur_extra, key))?;
+                Ok((next_state, next_extra))
             }
-            TravelEvent::Pop(tn) => {
-                cur_node_id = queue.pop_front().unwrap();
-                callback(TravelEvent::Pop((self.sam.get_state(cur_node_id), tn)))?;
-                Ok(())
+            TravelEvent::Pop(cur_tn, (cur_state, extra)) => {
+                let res = callback(TravelEvent::Pop((&cur_state, cur_tn), extra))?;
+                Ok((cur_state, res))
             }
-        })
+        }
     }
 
     pub fn dfs_along<
         TN: TrieNodeAlike<InnerType = T> + Clone,
-        E,
-        F: FnMut(TravelEvent<(GeneralSAMState<'_, T>, &TN), TN::InnerType>) -> Result<(), E>,
+        ExtraType,
+        ErrorType,
+        F: FnMut(
+            TravelEvent<(&GeneralSAMState<'_, T>, &TN), ErrorType, TN::InnerType>,
+        ) -> Result<ErrorType, ExtraType>,
     >(
         &self,
         trie_node: TN,
-        mut callback: F,
-    ) -> Result<(), E> {
-        let mut stack: Vec<usize> = Vec::new();
-
-        trie_node.dfs_travel(|event| match event {
-            TravelEvent::Push(tn, Some(key)) => {
-                let next_node_id = self
-                    .sam
-                    .node_pool
-                    .get(*stack.last().unwrap())
-                    .and_then(|x| x.trans.get(&key).copied())
-                    .unwrap_or(SAM_NIL_NODE_ID);
-                callback(TravelEvent::Push(
-                    (self.sam.get_state(next_node_id), tn),
-                    Some(key),
-                ))?;
-                stack.push(next_node_id);
-                Ok(())
-            }
-            TravelEvent::Push(tn, None) => {
-                callback(TravelEvent::Push(
-                    (self.sam.get_state(self.node_id), tn),
-                    None,
-                ))?;
-                stack.push(self.node_id);
-                Ok(())
-            }
-            TravelEvent::Pop(tn) => {
-                let node_id = stack.pop().unwrap();
-                callback(TravelEvent::Pop((self.sam.get_state(node_id), tn)))?;
-                Ok(())
-            }
-        })
+        callback: F,
+    ) -> Result<(), ExtraType> {
+        trie_node.dfs_travel(self.wrap_travel_along_callback(callback))
+    }
+
+    pub fn bfs_along<
+        TN: TrieNodeAlike<InnerType = T> + Clone,
+        ExtraType,
+        ErrorType,
+        F: FnMut(
+            TravelEvent<(&GeneralSAMState<'_, T>, &TN), ErrorType, TN::InnerType>,
+        ) -> Result<ErrorType, ExtraType>,
+    >(
+        &self,
+        trie_node: TN,
+        callback: F,
+    ) -> Result<(), ExtraType> {
+        trie_node.bfs_travel(self.wrap_travel_along_callback(callback))
     }
 }
diff --git a/src/trie_alike.rs b/src/trie_alike.rs
@@ -1,8 +1,9 @@
 use std::collections::VecDeque;
 
-pub enum TravelEvent<NodeType, KeyType> {
-    Push(NodeType, Option<KeyType>),
-    Pop(NodeType),
+pub enum TravelEvent<'s, NodeType, ExtraType, KeyType> {
+    PushRoot(NodeType),
+    Push(NodeType, &'s ExtraType, KeyType),
+    Pop(NodeType, ExtraType),
 }
 
 /// This trait provides the essential interfaces required by `GeneralSAM`
@@ -13,45 +14,59 @@ pub trait TrieNodeAlike {
     fn is_accepting(&self) -> bool;
     fn next_states(self) -> Self::NextStateIter;
 
-    fn bfs_travel<E, F: FnMut(TravelEvent<&Self, Self::InnerType>) -> Result<(), E>>(
+    fn bfs_travel<
+        ErrorType,
+        ExtraType,
+        F: FnMut(TravelEvent<&Self, ExtraType, Self::InnerType>) -> Result<ExtraType, ErrorType>,
+    >(
         self,
         mut callback: F,
-    ) -> Result<(), E>
+    ) -> Result<(), ErrorType>
     where
         Self: Sized,
     {
         let mut queue = VecDeque::new();
-        callback(TravelEvent::Push(&self, None))?;
-        queue.push_back(self);
-        while let Some(state) = queue.pop_front() {
-            callback(TravelEvent::Pop(&state))?;
+
+        let extra = callback(TravelEvent::PushRoot(&self))?;
+        queue.push_back((self, extra));
+
+        while let Some((state, cur_extra)) = queue.pop_front() {
+            let cur_extra = callback(TravelEvent::Pop(&state, cur_extra))?;
+
             for (t, v) in state.next_states() {
-                callback(TravelEvent::Push(&v, Some(t)))?;
-                queue.push_back(v);
+                let next_extra = callback(TravelEvent::Push(&v, &cur_extra, t))?;
+                queue.push_back((v, next_extra));
             }
         }
         Ok(())
     }
 
-    fn dfs_travel<E, F: FnMut(TravelEvent<&Self, Self::InnerType>) -> Result<(), E>>(
+    fn dfs_travel<
+        ErrorType,
+        ExtraType,
+        F: FnMut(TravelEvent<&Self, ExtraType, Self::InnerType>) -> Result<ExtraType, ErrorType>,
+    >(
         self,
         mut callback: F,
-    ) -> Result<(), E>
+    ) -> Result<(), ErrorType>
     where
         Self: Clone,
     {
         let mut stack = Vec::new();
 
-        callback(TravelEvent::Push(&self, None))?;
-        stack.push((self.clone(), self.next_states()));
+        let extra = callback(TravelEvent::PushRoot(&self))?;
+        stack.push((self.clone(), self.next_states(), extra));
 
-        while let Some((ref cur, ref mut iter)) = stack.last_mut() {
-            if let Some((key, next_state)) = iter.next() {
-                callback(TravelEvent::Push(&next_state, Some(key)))?;
-                stack.push((next_state.clone(), next_state.next_states()));
-            } else {
-                callback(TravelEvent::Pop(cur))?;
-                stack.pop();
+        while !stack.is_empty() {
+            if let Some((_, iter, extra)) = stack.last_mut() {
+                if let Some((key, next_state)) = iter.next() {
+                    let new_extra = callback(TravelEvent::Push(&next_state, extra, key))?;
+                    stack.push((next_state.clone(), next_state.next_states(), new_extra));
+                    continue;
+                }
+            }
+            if let Some((cur, _, extra)) = stack.pop() {
+                callback(TravelEvent::Pop(&cur, extra))?;
             }
         }
         Ok(())