custom fst automatons

milli/src/search/fst_utils.rs

@@ -0,0 +1,187 @@
+/// This mod is necessary until https://github.com/BurntSushi/fst/pull/137 gets merged.
+/// All credits for this code go to BurntSushi.
+use fst::Automaton;
+
+pub struct StartsWith<A>(pub A);
+
+/// The `Automaton` state for `StartsWith<A>`.
+pub struct StartsWithState<A: Automaton>(pub StartsWithStateKind<A>);
+
+impl<A: Automaton> Clone for StartsWithState<A>
+where
+    A::State: Clone,
+{
+    fn clone(&self) -> Self {
+        Self(self.0.clone())
+    }
+}
+
+/// The inner state of a `StartsWithState<A>`.
+pub enum StartsWithStateKind<A: Automaton> {
+    /// Sink state that is reached when the automaton has matched the prefix.
+    Done,
+    /// State in which the automaton is while it hasn't matched the prefix.
+    Running(A::State),
+}
+
+impl<A: Automaton> Clone for StartsWithStateKind<A>
+where
+    A::State: Clone,
+{
+    fn clone(&self) -> Self {
+        match self {
+            StartsWithStateKind::Done => StartsWithStateKind::Done,
+            StartsWithStateKind::Running(inner) => StartsWithStateKind::Running(inner.clone()),
+        }
+    }
+}
+
+impl<A: Automaton> Automaton for StartsWith<A> {
+    type State = StartsWithState<A>;
+
+    fn start(&self) -> StartsWithState<A> {
+        StartsWithState({
+            let inner = self.0.start();
+            if self.0.is_match(&inner) {
+                StartsWithStateKind::Done
+            } else {
+                StartsWithStateKind::Running(inner)
+            }
+        })
+    }
+    fn is_match(&self, state: &StartsWithState<A>) -> bool {
+        match state.0 {
+            StartsWithStateKind::Done => true,
+            StartsWithStateKind::Running(_) => false,
+        }
+    }
+    fn can_match(&self, state: &StartsWithState<A>) -> bool {
+        match state.0 {
+            StartsWithStateKind::Done => true,
+            StartsWithStateKind::Running(ref inner) => self.0.can_match(inner),
+        }
+    }
+    fn will_always_match(&self, state: &StartsWithState<A>) -> bool {
+        match state.0 {
+            StartsWithStateKind::Done => true,
+            StartsWithStateKind::Running(_) => false,
+        }
+    }
+    fn accept(&self, state: &StartsWithState<A>, byte: u8) -> StartsWithState<A> {
+        StartsWithState(match state.0 {
+            StartsWithStateKind::Done => StartsWithStateKind::Done,
+            StartsWithStateKind::Running(ref inner) => {
+                let next_inner = self.0.accept(inner, byte);
+                if self.0.is_match(&next_inner) {
+                    StartsWithStateKind::Done
+                } else {
+                    StartsWithStateKind::Running(next_inner)
+                }
+            }
+        })
+    }
+}
+/// An automaton that matches when one of its component automata match.
+#[derive(Clone, Debug)]
+pub struct Union<A, B>(pub A, pub B);
+
+/// The `Automaton` state for `Union<A, B>`.
+pub struct UnionState<A: Automaton, B: Automaton>(pub A::State, pub B::State);
+
+impl<A: Automaton, B: Automaton> Clone for UnionState<A, B>
+where
+    A::State: Clone,
+    B::State: Clone,
+{
+    fn clone(&self) -> Self {
+        Self(self.0.clone(), self.1.clone())
+    }
+}
+
+impl<A: Automaton, B: Automaton> Automaton for Union<A, B> {
+    type State = UnionState<A, B>;
+    fn start(&self) -> UnionState<A, B> {
+        UnionState(self.0.start(), self.1.start())
+    }
+    fn is_match(&self, state: &UnionState<A, B>) -> bool {
+        self.0.is_match(&state.0) || self.1.is_match(&state.1)
+    }
+    fn can_match(&self, state: &UnionState<A, B>) -> bool {
+        self.0.can_match(&state.0) || self.1.can_match(&state.1)
+    }
+    fn will_always_match(&self, state: &UnionState<A, B>) -> bool {
+        self.0.will_always_match(&state.0) || self.1.will_always_match(&state.1)
+    }
+    fn accept(&self, state: &UnionState<A, B>, byte: u8) -> UnionState<A, B> {
+        UnionState(self.0.accept(&state.0, byte), self.1.accept(&state.1, byte))
+    }
+}
+/// An automaton that matches when both of its component automata match.
+#[derive(Clone, Debug)]
+pub struct Intersection<A, B>(pub A, pub B);
+
+/// The `Automaton` state for `Intersection<A, B>`.
+pub struct IntersectionState<A: Automaton, B: Automaton>(pub A::State, pub B::State);
+
+impl<A: Automaton, B: Automaton> Clone for IntersectionState<A, B>
+where
+    A::State: Clone,
+    B::State: Clone,
+{
+    fn clone(&self) -> Self {
+        Self(self.0.clone(), self.1.clone())
+    }
+}
+
+impl<A: Automaton, B: Automaton> Automaton for Intersection<A, B> {
+    type State = IntersectionState<A, B>;
+    fn start(&self) -> IntersectionState<A, B> {
+        IntersectionState(self.0.start(), self.1.start())
+    }
+    fn is_match(&self, state: &IntersectionState<A, B>) -> bool {
+        self.0.is_match(&state.0) && self.1.is_match(&state.1)
+    }
+    fn can_match(&self, state: &IntersectionState<A, B>) -> bool {
+        self.0.can_match(&state.0) && self.1.can_match(&state.1)
+    }
+    fn will_always_match(&self, state: &IntersectionState<A, B>) -> bool {
+        self.0.will_always_match(&state.0) && self.1.will_always_match(&state.1)
+    }
+    fn accept(&self, state: &IntersectionState<A, B>, byte: u8) -> IntersectionState<A, B> {
+        IntersectionState(self.0.accept(&state.0, byte), self.1.accept(&state.1, byte))
+    }
+}
+/// An automaton that matches exactly when the automaton it wraps does not.
+#[derive(Clone, Debug)]
+pub struct Complement<A>(pub A);
+
+/// The `Automaton` state for `Complement<A>`.
+pub struct ComplementState<A: Automaton>(pub A::State);
+
+impl<A: Automaton> Clone for ComplementState<A>
+where
+    A::State: Clone,
+{
+    fn clone(&self) -> Self {
+        Self(self.0.clone())
+    }
+}
+
+impl<A: Automaton> Automaton for Complement<A> {
+    type State = ComplementState<A>;
+    fn start(&self) -> ComplementState<A> {
+        ComplementState(self.0.start())
+    }
+    fn is_match(&self, state: &ComplementState<A>) -> bool {
+        !self.0.is_match(&state.0)
+    }
+    fn can_match(&self, state: &ComplementState<A>) -> bool {
+        !self.0.will_always_match(&state.0)
+    }
+    fn will_always_match(&self, state: &ComplementState<A>) -> bool {
+        !self.0.can_match(&state.0)
+    }
+    fn accept(&self, state: &ComplementState<A>, byte: u8) -> ComplementState<A> {
+        ComplementState(self.0.accept(&state.0, byte))
+    }
+}