Generalize the NFA compiler info and specialize the PikeVM

Author: shilangyu

regex-automata/src/nfa/thompson/builder.rs

@@ -5,7 +5,7 @@ use alloc::{sync::Arc, vec, vec::Vec};
 use crate::{
     nfa::thompson::{
         error::BuildError,
-        nfa::{self, SparseTransitions, Transition, NFA},
+        nfa::{self, LookBehindInfo, SparseTransitions, Transition, NFA},
     },
     util::{
         look::{Look, LookMatcher},
@@ -340,12 +340,11 @@ pub struct Builder {
     /// contains a single regex, then `start_pattern[0]` and `start_anchored`
     /// are always equivalent.
     start_pattern: Vec<StateID>,
-    /// The starting states for each individual look-behind sub-expression.
-    start_look_behind: Vec<StateID>,
-    /// Among all look-behinds, this is the furthest offset (in bytes) from
-    /// the beginning of the main regex that a look-behind starts at.
-    /// If `None`, the offset is unbounded.
-    maximum_lookbehind_offset_from_start: Option<usize>,
+    /// A vector of meta-data information about each look-behind in this NFA.
+    ///
+    /// Must be stored in a depth-first pre-order with regards to the nesting
+    /// of look-behinds.
+    lookbehinds: Vec<LookBehindInfo>,
     /// A map from pattern ID to capture group index to name. (If no name
     /// exists, then a None entry is present. Thus, all capturing groups are
     /// present in this mapping.)
@@ -378,10 +377,7 @@ pub struct Builder {
 impl Builder {
     /// Create a new builder for hand-assembling NFAs.
     pub fn new() -> Builder {
-        Builder {
-            maximum_lookbehind_offset_from_start: Some(0),
-            ..Builder::default()
-        }
+        Builder::default()
     }
 
     /// Clear this builder.
@@ -394,7 +390,7 @@ impl Builder {
         self.pattern_id = None;
         self.states.clear();
         self.start_pattern.clear();
-        self.start_look_behind.clear();
+        self.lookbehinds.clear();
         self.captures.clear();
         self.memory_states = 0;
     }
@@ -459,10 +455,7 @@ impl Builder {
         remap.resize(self.states.len(), StateID::ZERO);
 
         nfa.set_starts(start_anchored, start_unanchored, &self.start_pattern);
-        nfa.set_look_behind_starts(self.start_look_behind.as_slice());
-        nfa.set_maximum_lookbehind_offset_from_start(
-            self.maximum_lookbehind_offset_from_start,
-        );
+        nfa.set_lookbehinds(self.lookbehinds.as_slice());
         nfa.set_captures(&self.captures).map_err(BuildError::captures)?;
         // The idea here is to convert our intermediate states to their final
         // form. The only real complexity here is the process of converting
@@ -724,21 +717,16 @@ impl Builder {
     /// running look-behind expressions. Additionally registers the furthest
     /// offset (in bytes) from the start of the main regex this look-behind
     /// starts.
-    pub fn start_look_behind(
+    ///
+    /// Look-behinds must be started in a depth-first pre-order fashion with
+    /// regards to the nesting of look-behinds.
+    pub fn start_lookbehind(
         &mut self,
         start_id: StateID,
         offset_from_start: Option<usize>,
     ) {
-        self.start_look_behind.push(start_id);
-
-        self.maximum_lookbehind_offset_from_start = match (
-            self.maximum_lookbehind_offset_from_start,
-            offset_from_start,
-        ) {
-            (Some(l1), Some(l2)) => Some(usize::max(l1, l2)),
-            // A None subsumes the entire result.
-            (None, _) | (_, None) => None,
-        };
+        self.lookbehinds
+            .push(LookBehindInfo::new(start_id, offset_from_start));
     }
 
     /// Add an "empty" NFA state.

regex-automata/src/nfa/thompson/compiler.rs

@@ -1060,7 +1060,7 @@ impl Compiler {
             self.c_at_least(&Hir::dot(hir::Dot::AnyByte), false, 0)?;
         self.builder
             .borrow_mut()
-            .start_look_behind(unanchored.start, start_offset);
+            .start_lookbehind(unanchored.start, start_offset);
 
         // When compiling the subexpression we temporarily change the starting
         // offset and restore it after. This way, the subexpression is relativized

regex-automata/src/nfa/thompson/nfa.rs

@@ -1108,15 +1108,8 @@ impl NFA {
 
     /// Returns the starting states for initializing look-behind evaluation.
     #[inline]
-    pub fn look_behind_starts(&self) -> &Vec<StateID> {
-        &self.0.start_look_behind
-    }
-
-    /// Returns the length (in bytes) of the longest string matched by any
-    /// look-behind sub-expression. If `None`, the length is unbounded.
-    #[inline]
-    pub fn maximum_lookbehind_offset_from_start(&self) -> Option<usize> {
-        self.0.maximum_lookbehind_offset_from_start
+    pub fn lookbehinds(&self) -> &Vec<LookBehindInfo> {
+        &self.0.lookbehinds
     }
 
     // FIXME: The `look_set_prefix_all` computation was not correct, and it
@@ -1283,19 +1276,48 @@ pub(super) struct Inner {
     /// This is needed to initialize the table for storing the result of
     /// look-around evaluation.
     lookaround_count: usize,
-    /// Contains the start states for each of the look-behind subexpressions.
-    start_look_behind: Vec<StateID>,
-    /// Among all look-behinds, this is the furthest offset (in bytes) from
-    /// the beginning of the main regex that a look-behind starts at.
-    /// If `None`, the offset is unbounded.
-    maximum_lookbehind_offset_from_start: Option<usize>,
+    /// A vector of meta-data information about each look-behind in this NFA.
+    ///
+    /// Must be stored in a depth-first pre-order with regards to the nesting
+    /// of look-behinds.
+    lookbehinds: Vec<LookBehindInfo>,
     /// Heap memory used indirectly by NFA states and other things (like the
     /// various capturing group representations above). Since each state
     /// might use a different amount of heap, we need to keep track of this
     /// incrementally.
     memory_extra: usize,
 }
 
+/// Information about a look-behind needed for execution.
+#[derive(Clone, Copy, Debug)]
+pub struct LookBehindInfo {
+    /// The id of the start state of the look-behind subexpression.
+    start_id: StateID,
+    /// The offset (in bytes) from the beginning of the main regex that a
+    /// look-behind starts at. If `None`, the offset is unbounded.
+    offset_from_start: Option<usize>,
+}
+
+impl LookBehindInfo {
+    pub(super) fn new(
+        start_id: StateID,
+        offset_from_start: Option<usize>,
+    ) -> Self {
+        Self { start_id, offset_from_start }
+    }
+
+    /// Start states of the look-behind subexpression.
+    pub(super) fn start_state(&self) -> StateID {
+        self.start_id
+    }
+
+    /// The offset (in bytes) from the beginning of the main regex that a
+    /// look-behind starts at. If `None`, the offset is unbounded.
+    pub(super) fn offset_from_start(&self) -> Option<usize> {
+        self.offset_from_start
+    }
+}
+
 impl Inner {
     /// Runs any last finalization bits and turns this into a full NFA.
     pub(super) fn into_nfa(mut self) -> NFA {
@@ -1438,19 +1460,12 @@ impl Inner {
         self.start_pattern = start_pattern.to_vec();
     }
 
-    pub(super) fn set_look_behind_starts(
-        &mut self,
-        look_behind_starts: &[StateID],
-    ) {
-        self.start_look_behind = look_behind_starts.to_vec();
-    }
-
-    pub(super) fn set_maximum_lookbehind_offset_from_start(
-        &mut self,
-        maximum_lookbehind_offset_from_start: Option<usize>,
-    ) {
-        self.maximum_lookbehind_offset_from_start =
-            maximum_lookbehind_offset_from_start;
+    /// Sets the look-behind information of this NFA.
+    ///
+    /// The slice must be in a depth-first pre-order with regards to the
+    /// nesting of look-behinds.
+    pub(super) fn set_lookbehinds(&mut self, lookbehinds: &[LookBehindInfo]) {
+        self.lookbehinds = lookbehinds.to_vec();
     }
 
     /// Sets the UTF-8 mode of this NFA.
@@ -1506,7 +1521,8 @@ impl Inner {
         for id in self.start_pattern.iter_mut() {
             *id = old_to_new[*id];
         }
-        for id in self.start_look_behind.iter_mut() {
+        for LookBehindInfo { start_id: id, .. } in self.lookbehinds.iter_mut()
+        {
             *id = old_to_new[*id];
         }
     }
@@ -1520,7 +1536,7 @@ impl fmt::Debug for Inner {
                 '^'
             } else if sid == self.start_unanchored {
                 '>'
-            } else if self.start_look_behind.contains(&sid) {
+            } else if self.lookbehinds.iter().any(|i| i.start_state() == sid) {
                 '<'
             } else {
                 ' '

regex-automata/src/nfa/thompson/pikevm.rs

@@ -278,7 +278,24 @@ impl Builder {
     /// given here is already built.
     pub fn build_from_nfa(&self, nfa: NFA) -> Result<PikeVM, BuildError> {
         nfa.look_set_any().available().map_err(BuildError::word)?;
-        Ok(PikeVM { config: self.config.clone(), nfa })
+
+        // The reverse of a depth-first pre-order is the depth-first
+        // reverse-post-order. This means, a look-around is always before its
+        // surrounding look-behinds in this vector.
+        let lookbehind_starts =
+            nfa.lookbehinds().iter().map(|i| i.start_state()).rev().collect();
+
+        let maximum_lookbehind_offset_from_start =
+            nfa.lookbehinds().iter().try_fold(0, |acc, curr| {
+                curr.offset_from_start().map(|o| usize::max(acc, o))
+            });
+
+        Ok(PikeVM {
+            config: self.config.clone(),
+            nfa,
+            lookbehind_starts,
+            maximum_lookbehind_offset_from_start,
+        })
     }
 
     /// Apply the given `PikeVM` configuration options to this builder.
@@ -387,6 +404,13 @@ impl Builder {
 pub struct PikeVM {
     config: Config,
     nfa: NFA,
+    /// Stored depth-first reverse-post-order with regards to the nesting
+    /// of look-behinds.
+    lookbehind_starts: Vec<StateID>,
+    /// Among all look-behinds, this is the furthest offset (in bytes) from
+    /// the beginning of the main regex that a look-behind starts at.
+    /// If `None`, the offset is unbounded.
+    maximum_lookbehind_offset_from_start: Option<usize>,
 }
 
 impl PikeVM {
@@ -1258,9 +1282,6 @@ impl PikeVM {
             Some(config) => config,
         };
 
-        let maximum_lookbehind_offset_from_start =
-            self.nfa.maximum_lookbehind_offset_from_start();
-
         let pre =
             if anchored { None } else { self.get_config().get_prefilter() };
         let Cache {
@@ -1281,18 +1302,17 @@ impl PikeVM {
             // start from 0.
             let start_position = usize::saturating_sub(
                 input.start(),
-                maximum_lookbehind_offset_from_start.unwrap_or(input.start()),
+                self.maximum_lookbehind_offset_from_start
+                    .unwrap_or(input.start()),
             );
 
             // This initializes the look-behind threads from the `start_position`
             // Note: since capture groups are not allowed inside look-behinds,
             // there won't be any Capture epsilon transitions and hence it is ok to
-            // use &mut [] for the slots parameter. We need to add the start states
-            // in reverse because more deeply nested look-behinds have a higher index
-            // but must be executed first, so that the result is available for the
-            // outer expression.
-            for look_behind_start in self.nfa.look_behind_starts().iter().rev()
-            {
+            // use &mut [] for the slots parameter. Since the start states are stored
+            // in depth-first reverse-post-order, more deeply nested look-behinds are
+            // executed first, so that the result is available for the outer expression.
+            for look_behind_start in &self.lookbehind_starts {
                 self.epsilon_closure(
                     stack,
                     &mut [],
@@ -1363,19 +1383,16 @@ impl PikeVM {
                                     at,
                                     usize::saturating_sub(
                                         span.start,
-                                        maximum_lookbehind_offset_from_start
+                                        self.maximum_lookbehind_offset_from_start
                                             .unwrap_or(span.start),
                                     ),
                                 );
                                 // If we resume from later than `at`, we need
                                 // to reinitialize the look-behind threads.
                                 if start_position != at {
                                     curr_lookaround.set.clear();
-                                    for look_behind_start in self
-                                        .nfa
-                                        .look_behind_starts()
-                                        .iter()
-                                        .rev()
+                                    for look_behind_start in
+                                        &self.lookbehind_starts
                                     {
                                         self.epsilon_closure(
                                             stack,
@@ -1598,7 +1615,7 @@ impl PikeVM {
             match_lookaround: _,
         } = cache;
 
-        for look_behind_start in self.nfa.look_behind_starts().iter().rev() {
+        for look_behind_start in &self.lookbehind_starts {
             self.epsilon_closure(
                 stack,
                 &mut [],