Store look-behind offsets separately

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, LookBehindInfo, SparseTransitions, Transition, NFA},
+        nfa::{self, LookBehindTree, SparseTransitions, Transition, NFA},
     },
     util::{
         look::{Look, LookMatcher},
@@ -340,11 +340,9 @@ pub struct Builder {
     /// contains a single regex, then `start_pattern[0]` and `start_anchored`
     /// are always equivalent.
     start_pattern: Vec<StateID>,
-    /// 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 vector of look-behinds appearing in the regex. Order reflects the
+    /// order in the regex.
+    lookbehinds: Vec<LookBehindTree>,
     /// 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.)
@@ -719,14 +717,21 @@ impl Builder {
     /// starts.
     ///
     /// Look-behinds must be started in a depth-first pre-order fashion with
-    /// regards to the nesting of look-behinds.
+    /// regards to the nesting of look-behinds. The nesting path is stored
+    /// as indices in `path`.
     pub fn start_lookbehind(
         &mut self,
         start_id: StateID,
         offset_from_start: Option<usize>,
+        path: &[usize],
     ) {
-        self.lookbehinds
-            .push(LookBehindInfo::new(start_id, offset_from_start));
+        let mut current = &mut self.lookbehinds;
+
+        for index in path {
+            current = current[*index].children_mut();
+        }
+
+        current.push(LookBehindTree::new(start_id, offset_from_start));
     }
 
     /// Add an "empty" NFA state.

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

@@ -718,6 +718,8 @@ pub struct Compiler {
     /// the current look-behind expression. When `None`, the distance can be
     /// seen as infinity.
     current_lookbehind_offset_from_start: RefCell<Option<usize>>,
+    /// The current path of look-behind nesting.
+    lookbehind_nesting_path: RefCell<Vec<usize>>,
 }
 
 impl Compiler {
@@ -732,6 +734,7 @@ impl Compiler {
             utf8_suffix: RefCell::new(Utf8SuffixMap::new(1000)),
             lookaround_index: RefCell::new(SmallIndex::ZERO),
             current_lookbehind_offset_from_start: RefCell::new(Some(0)),
+            lookbehind_nesting_path: RefCell::new(vec![0]),
         }
     }
 
@@ -971,6 +974,9 @@ impl Compiler {
         self.builder
             .borrow_mut()
             .set_size_limit(self.config.get_nfa_size_limit())?;
+        *self.lookaround_index.borrow_mut() = SmallIndex::ZERO;
+        *self.lookbehind_nesting_path.borrow_mut() = vec![0];
+        *self.current_lookbehind_offset_from_start.borrow_mut() = Some(0);
 
         // We always add an unanchored prefix unless we were specifically told
         // not to (for tests only), or if we know that the regex is anchored
@@ -1058,15 +1064,22 @@ impl Compiler {
 
         let unanchored =
             self.c_at_least(&Hir::dot(hir::Dot::AnyByte), false, 0)?;
-        self.builder
-            .borrow_mut()
-            .start_lookbehind(unanchored.start, start_offset);
+        self.builder.borrow_mut().start_lookbehind(
+            unanchored.start,
+            start_offset,
+            self.lookbehind_nesting_path.borrow().split_last().unwrap().1,
+        );
 
         // When compiling the subexpression we temporarily change the starting
         // offset and restore it after. This way, the subexpression is relativized
-        // to our current offset.
+        // to our current offset. We also update the path to the current lookbehind
+        // expression.
+        self.lookbehind_nesting_path.borrow_mut().push(0);
         *self.current_lookbehind_offset_from_start.borrow_mut() = start_offset;
         let sub = self.c(lookaround.sub())?;
+        let mut path = self.lookbehind_nesting_path.borrow_mut();
+        path.pop();
+        *path.last_mut().unwrap() += 1;
         *self.current_lookbehind_offset_from_start.borrow_mut() =
             relative_start;
 

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

@@ -1108,7 +1108,7 @@ impl NFA {
 
     /// Returns the starting states for initializing look-behind evaluation.
     #[inline]
-    pub fn lookbehinds(&self) -> &Vec<LookBehindInfo> {
+    pub fn lookbehinds(&self) -> &[LookBehindTree] {
         &self.0.lookbehinds
     }
 
@@ -1276,46 +1276,83 @@ pub(super) struct Inner {
     /// This is needed to initialize the table for storing the result of
     /// look-around evaluation.
     lookaround_count: 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 vector of look-behinds appearing in the regex. Order reflects the
+    /// order in the regex.
+    lookbehinds: Vec<LookBehindTree>,
     /// 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.
+/// Information about a look-behinds needed for execution. It preserves the
+/// nesting structure of look-behinds.
+#[derive(Clone, Debug)]
+pub struct LookBehindTree {
     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>,
+    children: Vec<LookBehindTree>,
 }
 
-impl LookBehindInfo {
-    pub(super) fn new(
-        start_id: StateID,
-        offset_from_start: Option<usize>,
-    ) -> Self {
-        Self { start_id, offset_from_start }
+impl LookBehindTree {
+    pub fn new(start_id: StateID, offset_from_start: Option<usize>) -> Self {
+        Self { start_id, offset_from_start, children: Vec::new() }
     }
 
-    /// Start states of the look-behind subexpression.
-    pub(super) fn start_state(&self) -> StateID {
+    /// The id of the start state of the look-behind subexpression.
+    pub fn start_id(&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> {
+    pub fn offset_from_start(&self) -> Option<usize> {
         self.offset_from_start
     }
+
+    /// The look-behinds this look-behind contains. Order reflects the order
+    /// in the regex.
+    pub fn children(&self) -> &[LookBehindTree] {
+        &self.children
+    }
+
+    /// Calls `fun` on this look-behind tree and all of its children in pre-order.
+    /// `fun` should return `true` if the traversal should continue and `false`
+    /// if it should stop.
+    ///
+    /// The return value indicates whether the traversal was at any point stopped.
+    pub fn preorder(&self, fun: &impl Fn(&LookBehindTree) -> bool) -> bool {
+        if !fun(self) {
+            return false;
+        }
+        for child in &self.children {
+            if !child.preorder(fun) {
+                return false;
+            }
+        }
+        true
+    }
+
+    /// Like [`preorder`], but allows mutating the nodes.
+    pub fn preorder_mut(
+        &mut self,
+        fun: &impl Fn(&mut LookBehindTree) -> bool,
+    ) -> bool {
+        if !fun(self) {
+            return false;
+        }
+        for child in &mut self.children {
+            if !child.preorder_mut(fun) {
+                return false;
+            }
+        }
+        true
+    }
+
+    pub fn children_mut(&mut self) -> &mut Vec<LookBehindTree> {
+        &mut self.children
+    }
 }
 
 impl Inner {
@@ -1464,7 +1501,7 @@ impl Inner {
     ///
     /// 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]) {
+    pub(super) fn set_lookbehinds(&mut self, lookbehinds: &[LookBehindTree]) {
         self.lookbehinds = lookbehinds.to_vec();
     }
 
@@ -1521,9 +1558,12 @@ impl Inner {
         for id in self.start_pattern.iter_mut() {
             *id = old_to_new[*id];
         }
-        for LookBehindInfo { start_id: id, .. } in self.lookbehinds.iter_mut()
-        {
-            *id = old_to_new[*id];
+
+        for lbs in self.lookbehinds.iter_mut() {
+            lbs.preorder_mut(&|e| {
+                e.start_id = old_to_new[e.start_id];
+                true
+            });
         }
     }
 }
@@ -1536,7 +1576,11 @@ impl fmt::Debug for Inner {
                 '^'
             } else if sid == self.start_unanchored {
                 '>'
-            } else if self.lookbehinds.iter().any(|i| i.start_state() == sid) {
+            } else if self
+                .lookbehinds
+                .iter()
+                .any(|i| !i.preorder(&|e| e.start_id() != sid))
+            {
                 '<'
             } else {
                 ' '