unix-ffi/re: Cache the compiled patterns.

Every call to re.search(), and to the functions next to it, compiled the
pattern it was given.  Keep the compiled patterns in a small cache
instead, the way CPython does, so that using the same pattern again does
not compile it a second time.  compile() returns the cached pattern as
well, so re.compile(p) is re.compile(p), as it is in CPython.  Matching
against a repeated pattern gets about twice as fast, compiling one about
eight times.

Because MicroPython cannot release a compiled pattern by itself, the
cache also decides what is kept: a cached pattern stays for the lifetime
of the program, and a pattern that this module compiled for its own use
is freed again afterwards.

The cache owns what it holds and never evicts it.  A pattern that is
still in use, by the caller or by a call further up the stack, must not
be freed underneath it, which a replacement callback passed to sub() can
otherwise trigger.  The cache is bounded instead: once it is full,
further patterns are compiled and, where this module owns them, freed
again after use.

Signed-off-by: Kirill Lukonin (Evil Wireless Man) <klukonin@gmail.com>

Author: klukonin

unix-ffi/re/re.py

@@ -91,12 +91,18 @@ def span(self, n=0):
 class PCREPattern:
     def __init__(self, compiled_ptn):
         self.obj = compiled_ptn
+        self.key = None  # set while this pattern is held by the cache
 
     def _free(self):
         # MicroPython does not run __del__ on instances of Python classes, so
         # the compiled pattern cannot be released by the garbage collector and
         # has to be freed explicitly.
         if self.obj is not None:
+            if self.key is not None:
+                # Drop the pattern from the cache first, so that nothing hands
+                # out a pointer that is about to become invalid.
+                del _cache[self.key]
+                self.key = None
             pcre2_code_free(self.obj)
             self.obj = None
 
@@ -184,7 +190,7 @@ def findall(self, s):
             start = end
 
 
-def compile(pattern, flags=0):
+def _compile(pattern, flags):
     # These are output arguments and must be writable and of the size that
     # pcre2_compile() writes: int for the error code, PCRE2_SIZE for the offset.
     errcode = array.array("i", [0])
@@ -194,48 +200,82 @@ def compile(pattern, flags=0):
     return PCREPattern(regex)
 
 
-# The functions below compile a pattern that is not visible to the caller, so
-# they must free it again.  The match objects they return do not refer to it.
+# Compiled patterns are cached, the way CPython does it, so that using the same
+# pattern again does not compile it a second time.  compile() returns the
+# cached pattern, so re.compile(p) is re.compile(p), as in CPython.
+#
+# The cache owns the patterns it holds and never evicts them.  A pattern that
+# is still being used, either by the caller or by a call further up the stack,
+# must not be freed underneath it; a replacement callback passed to sub() can
+# otherwise trigger exactly that.  The cache is bounded instead: once it is
+# full, further patterns are compiled and, where this module owns them, freed
+# again after use.
+_MAXCACHE = 32
+_cache = {}
+
+
+def _cached(pattern, flags):
+    # Return the compiled pattern, and whether the caller has to free it.
+    key = (pattern, flags)
+    r = _cache.get(key)
+    if r is not None:
+        return r, False
+    r = _compile(pattern, flags)
+    if len(_cache) < _MAXCACHE:
+        _cache[key] = r
+        r.key = key
+        return r, False
+    return r, True
+
+
+def compile(pattern, flags=0):
+    # The pattern belongs to the caller, so it is never freed here.
+    return _cached(pattern, flags)[0]
 
 
 def search(pattern, string, flags=0):
-    r = compile(pattern, flags)
+    r, owned = _cached(pattern, flags)
     try:
         return r.search(string)
     finally:
-        r._free()
+        if owned:
+            r._free()
 
 
 def match(pattern, string, flags=0):
-    r = compile(pattern, flags | PCRE2_ANCHORED)
+    r, owned = _cached(pattern, flags | PCRE2_ANCHORED)
     try:
         return r.search(string)
     finally:
-        r._free()
+        if owned:
+            r._free()
 
 
 def sub(pattern, repl, s, count=0, flags=0):
-    r = compile(pattern, flags)
+    r, owned = _cached(pattern, flags)
     try:
         return r.sub(repl, s, count)
     finally:
-        r._free()
+        if owned:
+            r._free()
 
 
 def split(pattern, s, maxsplit=0, flags=0):
-    r = compile(pattern, flags)
+    r, owned = _cached(pattern, flags)
     try:
         return r.split(s, maxsplit)
     finally:
-        r._free()
+        if owned:
+            r._free()
 
 
 def findall(pattern, s, flags=0):
-    r = compile(pattern, flags)
+    r, owned = _cached(pattern, flags)
     try:
         return r.findall(s)
     finally:
-        r._free()
+        if owned:
+            r._free()
 
 
 def escape(s):

unix-ffi/re/test_re_leak.py

@@ -6,6 +6,10 @@
 # A pattern returned by re.compile() and kept by the caller is not covered
 # here.  MicroPython does not run __del__ on instances of Python classes, so
 # such a pattern can only be released explicitly.
+#
+# The bounded cache that the module level functions keep is covered: it must
+# not grow past its limit, and the patterns that do not fit into it must be
+# freed again.
 
 import gc
 import re
@@ -21,6 +25,7 @@ def rss():
     rss()
 except OSError:
     # No /proc, so memory use cannot be measured here.
+    print("SKIP")
     raise SystemExit
 
 
@@ -97,3 +102,48 @@ def many_groups():
 
 
 check_no_leak("pattern with several groups", many_groups)
+
+
+# compile() returns the cached pattern, the way CPython does, so compiling the
+# same pattern again does not allocate.
+assert re.compile("a(b)c") is re.compile("a(b)c")
+check_no_leak("re.compile() with the same pattern", lambda: re.compile("a(b)c"))
+
+# _free() drops the pattern from the cache, so that nothing afterwards hands
+# out a pointer to memory that has been released.
+r = re.compile("zz(y)")
+r._free()
+assert re.search("zz(y)", "xxzzyxx").group(0) == "zzy"
+
+
+# The module level functions cache the patterns they compile.  That cache must
+# stay bounded, and a pattern that does not fit into it has to be freed again.
+counter = [0]
+
+
+def distinct_patterns():
+    counter[0] += 1
+    re.search("a%dc" % counter[0], "xxabcxx")
+
+
+# Push far more distinct patterns through the cache than it can hold: it has
+# to stop growing.
+for _ in range(re._MAXCACHE * 4):
+    distinct_patterns()
+assert len(re._cache) <= re._MAXCACHE, len(re._cache)
+
+check_no_leak("re.search() with distinct patterns", distinct_patterns)
+assert len(re._cache) <= re._MAXCACHE, len(re._cache)
+
+
+# A replacement callback runs while sub() is still using its own pattern, and
+# may push further patterns through the cache.  The pattern that is in use must
+# survive that.
+def reentrant_repl(m):
+    counter[0] += 1
+    re.search("z%dz" % counter[0], "nothing here")
+    return "z"
+
+
+check_no_leak("re.sub() with a reentrant callback", lambda: re.sub("a", reentrant_repl, "caaab"))
+assert len(re._cache) <= re._MAXCACHE, len(re._cache)