automata: add new 'WhichCaptures' config

BurntSushi · BurntSushi · commit e10c9d7b56d3 · 2023-08-05T14:28:50.000-04:00
This is the first step in fixing a regression in memory usage. The
underlying problem is that regex-automata now natively supports
multi-pattern regexes *with* capturing support. Unfortunately though,
this overall doesn't work too well with the current design of the
PikeVM, because the amount of memory used is `len(captures) *
len(states)`. So basically, as the regex and number of captures
increases, the amount of memory used gets quite high.

This is new functionality that we hope to improve upon over time, so
it's not too big of a deal on its own. But it turns out this impacts
previous uses of RegexSet that have capture groups. The old
implementation just ignored these capture groups because they weren't
supported in a RegexSet, and thus there were no memory problems. But in
the new implementation, nothing tells it that it's okay to ignore the
capture groups. So it winds up allocating space for them even though the
RegexSet APIs don't provide any of that functionality.

So my plan to fix this is to introduce a new configuration knob for
controlling more granularly which capture states are compiled into the
NFA. Previously we only supported "all of them" or "none of them." This
commit adds a new (backwards compatible) knob that also permits "just
implicit groups." That is, one capture group per pattern. This hopefully
leads to less memory usage overall. (Well, it will certaintly be less,
but hopefully it's a big reduction.) We don't actually change anything
here. We just add a new `Config::which_captures` knob, implement the
existing `Config::captures` in terms of `Config::which_captures` and
deprecate `Config::captures`.

If this winds up not being sufficient, then we may need to adapt the
PikeVM to work without any capture groups at all and instead just report
which patterns match. Which is... probably fine?
diff --git a/regex-automata/src/dfa/dense.rs b/regex-automata/src/dfa/dense.rs
@@ -1170,7 +1170,10 @@ impl Builder {
             .clone()
             // We can always forcefully disable captures because DFAs do not
             // support them.
-            .configure(thompson::Config::new().captures(false))
+            .configure(
+                thompson::Config::new()
+                    .which_captures(thompson::WhichCaptures::None),
+            )
             .build_many(patterns)
             .map_err(BuildError::nfa)?;
         self.build_from_nfa(&nfa)
diff --git a/regex-automata/src/hybrid/dfa.rs b/regex-automata/src/hybrid/dfa.rs
@@ -3973,7 +3973,10 @@ impl Builder {
             .clone()
             // We can always forcefully disable captures because DFAs do not
             // support them.
-            .configure(thompson::Config::new().captures(false))
+            .configure(
+                thompson::Config::new()
+                    .which_captures(thompson::WhichCaptures::None),
+            )
             .build_many(patterns)
             .map_err(BuildError::nfa)?;
         self.build_from_nfa(nfa)
diff --git a/regex-automata/src/meta/strategy.rs b/regex-automata/src/meta/strategy.rs
@@ -13,7 +13,7 @@ use crate::{
         regex::{Cache, RegexInfo},
         reverse_inner, wrappers,
     },
-    nfa::thompson::{self, NFA},
+    nfa::thompson::{self, WhichCaptures, NFA},
     util::{
         captures::{Captures, GroupInfo},
         look::LookMatcher,
@@ -452,7 +452,7 @@ impl Core {
             .utf8(info.config().get_utf8_empty())
             .nfa_size_limit(info.config().get_nfa_size_limit())
             .shrink(false)
-            .captures(true)
+            .which_captures(WhichCaptures::All)
             .look_matcher(lookm);
         let nfa = thompson::Compiler::new()
             .configure(thompson_config.clone())
@@ -499,7 +499,10 @@ impl Core {
                     // useful with capturing groups in reverse. And of course,
                     // the lazy DFA ignores capturing groups in all cases.
                     .configure(
-                        thompson_config.clone().captures(false).reverse(true),
+                        thompson_config
+                            .clone()
+                            .which_captures(WhichCaptures::None)
+                            .reverse(true),
                     )
                     .build_many_from_hir(hirs)
                     .map_err(BuildError::nfa)?;
@@ -1480,7 +1483,7 @@ impl ReverseInner {
             .utf8(core.info.config().get_utf8_empty())
             .nfa_size_limit(core.info.config().get_nfa_size_limit())
             .shrink(false)
-            .captures(false)
+            .which_captures(WhichCaptures::None)
             .look_matcher(lookm);
         let result = thompson::Compiler::new()
             .configure(thompson_config)
diff --git a/regex-automata/src/nfa/thompson/compiler.rs b/regex-automata/src/nfa/thompson/compiler.rs
@@ -30,7 +30,7 @@ pub struct Config {
     reverse: Option<bool>,
     nfa_size_limit: Option<Option<usize>>,
     shrink: Option<bool>,
-    captures: Option<bool>,
+    which_captures: Option<WhichCaptures>,
     look_matcher: Option<LookMatcher>,
     #[cfg(test)]
     unanchored_prefix: Option<bool>,
@@ -178,12 +178,15 @@ impl Config {
     /// ```
     /// use regex_automata::{
     ///     dfa::{self, Automaton},
-    ///     nfa::thompson::NFA,
+    ///     nfa::thompson::{NFA, WhichCaptures},
     ///     HalfMatch, Input,
     /// };
     ///
     /// let dfa = dfa::dense::Builder::new()
-    ///     .thompson(NFA::config().captures(false).reverse(true))
+    ///     .thompson(NFA::config()
+    ///         .which_captures(WhichCaptures::None)
+    ///         .reverse(true)
+    ///     )
     ///     .build("baz[0-9]+")?;
     /// let expected = Some(HalfMatch::must(0, 3));
     /// assert_eq!(
@@ -277,10 +280,12 @@ impl Config {
     ///
     /// ```
     /// # if cfg!(miri) { return Ok(()); } // miri takes too long
-    /// use regex_automata::nfa::thompson::NFA;
+    /// use regex_automata::nfa::thompson::{NFA, WhichCaptures};
     ///
     /// // Currently we have to disable captures when enabling reverse NFA.
-    /// let config = NFA::config().captures(false).reverse(true);
+    /// let config = NFA::config()
+    ///     .which_captures(WhichCaptures::None)
+    ///     .reverse(true);
     /// let not_shrunk = NFA::compiler()
     ///     .configure(config.clone().shrink(false))
     ///     .build(r"\w")?;
@@ -314,18 +319,70 @@ impl Config {
     /// require capturing groups to be present in the NFA. Building a Pike VM
     /// with an NFA without capturing groups will result in an error.
     ///
+    /// (Note that since this method is deprecated, the example below uses
+    /// [`Config::which_captures`] to disable capture states.)
+    ///
     /// ```
-    /// use regex_automata::nfa::thompson::{pikevm::PikeVM, NFA};
+    /// use regex_automata::nfa::thompson::{
+    ///     pikevm::PikeVM,
+    ///     NFA,
+    ///     WhichCaptures,
+    /// };
     ///
     /// let nfa = NFA::compiler()
-    ///     .configure(NFA::config().captures(false))
+    ///     .configure(NFA::config().which_captures(WhichCaptures::None))
     ///     .build(r"[a-z]+")?;
     /// assert!(PikeVM::new_from_nfa(nfa).is_err());
     ///
     /// # Ok::<(), Box<dyn std::error::Error>>(())
     /// ```
-    pub fn captures(mut self, yes: bool) -> Config {
-        self.captures = Some(yes);
+    #[deprecated(since = "0.3.5", note = "use which_captures instead")]
+    pub fn captures(self, yes: bool) -> Config {
+        self.which_captures(if yes {
+            WhichCaptures::All
+        } else {
+            WhichCaptures::None
+        })
+    }
+
+    /// Configures what kinds of capture groups are compiled into
+    /// [`State::Capture`](crate::nfa::thompson::State::Capture) states in a
+    /// Thompson NFA.
+    ///
+    /// Currently, using any option except for [`WhichCaptures::None`] requires
+    /// disabling the [`reverse`](Config::reverse) setting. If both are
+    /// enabled, then the compiler will return an error. It is expected that
+    /// this limitation will be lifted in the future.
+    ///
+    /// This is set to [`WhichCaptures::All`] by default. Callers may wish to
+    /// use [`WhichCaptures::Implicit`] in cases where one wants avoid the
+    /// overhead of capture states for explicit groups. Usually this occurs
+    /// when one wants to use the `PikeVM` only for determining the overall
+    /// match. Otherwise, the `PikeVM` could use much more memory than is
+    /// necessary.
+    ///
+    /// # Example
+    ///
+    /// This example demonstrates that some regex engines, like the Pike VM,
+    /// require capturing groups to be present in the NFA. Building a Pike VM
+    /// with an NFA without capturing groups will result in an error.
+    ///
+    /// ```
+    /// use regex_automata::nfa::thompson::{
+    ///     pikevm::PikeVM,
+    ///     NFA,
+    ///     WhichCaptures,
+    /// };
+    ///
+    /// let nfa = NFA::compiler()
+    ///     .configure(NFA::config().which_captures(WhichCaptures::None))
+    ///     .build(r"[a-z]+")?;
+    /// assert!(PikeVM::new_from_nfa(nfa).is_err());
+    ///
+    /// # Ok::<(), Box<dyn std::error::Error>>(())
+    /// ```
+    pub fn which_captures(mut self, which_captures: WhichCaptures) -> Config {
+        self.which_captures = Some(which_captures);
         self
     }
 
@@ -405,8 +462,14 @@ impl Config {
     }
 
     /// Return whether NFA compilation is configured to produce capture states.
+    #[deprecated(since = "0.3.5", note = "use get_which_captures instead")]
     pub fn get_captures(&self) -> bool {
-        self.captures.unwrap_or(true)
+        self.get_which_captures().is_any()
+    }
+
+    /// Return what kinds of capture states will be compiled into an NFA.
+    pub fn get_which_captures(&self) -> WhichCaptures {
+        self.which_captures.unwrap_or(WhichCaptures::All)
     }
 
     /// Return the look-around matcher for this NFA.
@@ -439,14 +502,65 @@ impl Config {
             reverse: o.reverse.or(self.reverse),
             nfa_size_limit: o.nfa_size_limit.or(self.nfa_size_limit),
             shrink: o.shrink.or(self.shrink),
-            captures: o.captures.or(self.captures),
+            which_captures: o.which_captures.or(self.which_captures),
             look_matcher: o.look_matcher.or_else(|| self.look_matcher.clone()),
             #[cfg(test)]
             unanchored_prefix: o.unanchored_prefix.or(self.unanchored_prefix),
         }
     }
 }
 
+/// A configuration indicating which kinds of
+/// [`State::Capture`](crate::nfa::thompson::State::Capture) states to include.
+///
+/// This configuration can be used with [`Config::which_captures`] to control
+/// which capture states are compiled into a Thompson NFA.
+///
+/// The default configuration is [`WhichCaptures::All`].
+#[derive(Clone, Copy, Debug)]
+pub enum WhichCaptures {
+    /// All capture states, including those corresponding to both implicit and
+    /// explicit capture groups, are included in the Thompson NFA.
+    All,
+    /// Only capture states corresponding to implicit capture groups are
+    /// included. Implicit capture groups appear in every pattern implicitly
+    /// and correspond to the overall match of a pattern.
+    ///
+    /// This is useful when one only cares about the overall match of a
+    /// pattern. By excluding capture states from explicit capture groups,
+    /// one might be able to reduce the memory usage of a multi-pattern regex
+    /// substantially if it was otherwise written to have many explicit capture
+    /// groups.
+    Implicit,
+    /// No capture states are compiled into the Thompson NFA.
+    ///
+    /// This is useful when capture states are either not needed (for example,
+    /// if one is only trying to build a DFA) or if they aren't supported (for
+    /// example, a reverse NFA).
+    None,
+}
+
+impl Default for WhichCaptures {
+    fn default() -> WhichCaptures {
+        WhichCaptures::All
+    }
+}
+
+impl WhichCaptures {
+    /// Returns true if this configuration indicates that no capture states
+    /// should be produced in an NFA.
+    pub fn is_none(&self) -> bool {
+        matches!(*self, WhichCaptures::None)
+    }
+
+    /// Returns true if this configuration indicates that some capture states
+    /// should be added to an NFA. Note that this might only include capture
+    /// states for implicit capture groups.
+    pub fn is_any(&self) -> bool {
+        !self.is_none()
+    }
+}
+
 /*
 This compiler below uses Thompson's construction algorithm. The compiler takes
 a regex-syntax::Hir as input and emits an NFA graph as output. The NFA graph
@@ -800,7 +914,9 @@ impl Compiler {
         if exprs.len() > PatternID::LIMIT {
             return Err(BuildError::too_many_patterns(exprs.len()));
         }
-        if self.config.get_reverse() && self.config.get_captures() {
+        if self.config.get_reverse()
+            && self.config.get_which_captures().is_any()
+        {
             return Err(BuildError::unsupported_captures());
         }
 
@@ -978,7 +1094,7 @@ impl Compiler {
         name: Option<&str>,
         expr: &Hir,
     ) -> Result<ThompsonRef, BuildError> {
-        if !self.config.get_captures() {
+        if self.config.get_which_captures().is_none() {
             return self.c(expr);
         }
 
@@ -1728,9 +1844,15 @@ mod tests {
         util::primitives::{PatternID, StateID},
     };
 
+    use super::*;
+
     fn build(pattern: &str) -> NFA {
         NFA::compiler()
-            .configure(NFA::config().captures(false).unanchored_prefix(false))
+            .configure(
+                NFA::config()
+                    .which_captures(WhichCaptures::None)
+                    .unanchored_prefix(false),
+            )
             .build(pattern)
             .unwrap()
     }
@@ -1794,7 +1916,7 @@ mod tests {
     #[test]
     fn compile_unanchored_prefix() {
         let nfa = NFA::compiler()
-            .configure(NFA::config().captures(false))
+            .configure(NFA::config().which_captures(WhichCaptures::None))
             .build(r"a")
             .unwrap();
         assert_eq!(
@@ -1827,7 +1949,11 @@ mod tests {
 
         // Check that non-UTF-8 literals work.
         let nfa = NFA::compiler()
-            .configure(NFA::config().captures(false).unanchored_prefix(false))
+            .configure(
+                NFA::config()
+                    .which_captures(WhichCaptures::None)
+                    .unanchored_prefix(false),
+            )
             .syntax(crate::util::syntax::Config::new().utf8(false))
             .build(r"(?-u)\xFF")
             .unwrap();
@@ -1937,7 +2063,7 @@ mod tests {
         let nfa = NFA::compiler()
             .configure(
                 NFA::config()
-                    .captures(false)
+                    .which_captures(WhichCaptures::None)
                     .reverse(true)
                     .shrink(false)
                     .unanchored_prefix(false),
@@ -1965,7 +2091,11 @@ mod tests {
     #[test]
     fn compile_many_start_pattern() {
         let nfa = NFA::compiler()
-            .configure(NFA::config().captures(false).unanchored_prefix(false))
+            .configure(
+                NFA::config()
+                    .which_captures(WhichCaptures::None)
+                    .unanchored_prefix(false),
+            )
             .build_many(&["a", "b"])
             .unwrap();
         assert_eq!(
@@ -1993,7 +2123,9 @@ mod tests {
         use regex_syntax::hir::{Class, ClassBytes, Hir};
 
         let hir = Hir::class(Class::Bytes(ClassBytes::new(vec![])));
-        let config = NFA::config().captures(false).unanchored_prefix(false);
+        let config = NFA::config()
+            .which_captures(WhichCaptures::None)
+            .unanchored_prefix(false);
         let nfa =
             NFA::compiler().configure(config).build_from_hir(&hir).unwrap();
         assert_eq!(nfa.states(), &[s_fail(), s_match(0)]);
@@ -2005,7 +2137,9 @@ mod tests {
         use regex_syntax::hir::{Class, ClassUnicode, Hir};
 
         let hir = Hir::class(Class::Unicode(ClassUnicode::new(vec![])));
-        let config = NFA::config().captures(false).unanchored_prefix(false);
+        let config = NFA::config()
+            .which_captures(WhichCaptures::None)
+            .unanchored_prefix(false);
         let nfa =
             NFA::compiler().configure(config).build_from_hir(&hir).unwrap();
         assert_eq!(nfa.states(), &[s_fail(), s_match(0)]);
diff --git a/regex-automata/src/nfa/thompson/mod.rs b/regex-automata/src/nfa/thompson/mod.rs
@@ -78,4 +78,4 @@ pub use self::{
     },
 };
 #[cfg(feature = "syntax")]
-pub use compiler::{Compiler, Config};
+pub use compiler::{Compiler, Config, WhichCaptures};
diff --git a/regex-automata/src/nfa/thompson/nfa.rs b/regex-automata/src/nfa/thompson/nfa.rs
diff --git a/regex-automata/src/util/captures.rs b/regex-automata/src/util/captures.rs
