Make weakref memoization more test-friendly

Do pre-marking with a `set` method that we can patch in the tests.

Run assertions both before and after `.mark_deterministic()`.

Move _WeakrefMemoized to the end of the file as it's just noise for
Mypy, not the actually interesting code to read.

Author: blyxxyz

nnf/__init__.py

@@ -726,9 +726,9 @@ def to_MODS(self) -> 'Or[And[Var]]':
             And(Var(name, val) for name, val in model.items())
             for model in self.models()
         )
-        NNF.is_MODS.memo[new] = True
-        NNF._is_DNF_loose.memo[new] = True
-        NNF._is_DNF_strict.memo[new] = True
+        NNF.is_MODS.set(new, True)
+        NNF._is_DNF_loose.set(new, True)
+        NNF._is_DNF_strict.set(new, True)
         return new
 
     def to_model(self) -> Model:
@@ -810,7 +810,7 @@ def smooth(node: NNF) -> NNF:
             return new
 
         ret = smooth(self)
-        NNF.smooth.memo[ret] = True
+        NNF.smooth.set(ret, True)
         return ret
 
     def simplify(self, merge_nodes: bool = True) -> 'NNF':

nnf/dimacs.py

@@ -313,5 +313,5 @@ def _parse_cnf(tokens: t.Iterable[str]) -> And[Or[Var]]:
         # Adding an empty clause is not desirable
         clauses.add(Or(clause))
     sentence = And(clauses)
-    NNF._is_CNF_loose.memo[sentence] = True
+    NNF._is_CNF_loose.set(sentence, True)
     return sentence

nnf/dsharp.py

@@ -144,5 +144,5 @@ def compile(
 
     result = loads(out, var_labels=var_labels)
     result.mark_deterministic()
-    NNF.decomposable.memo[result] = True
+    NNF.decomposable.set(result, True)
     return result

nnf/tseitin.py

@@ -53,6 +53,6 @@ def process_node(node: NNF) -> Var:
     root = process_node(theory)
     clauses.append(Or({root}))
     ret = And(clauses)
-    NNF._is_CNF_loose.memo[ret] = True
-    NNF._is_CNF_strict.memo[ret] = True
+    NNF._is_CNF_loose.set(ret, True)
+    NNF._is_CNF_strict.set(ret, True)
     return ret

nnf/util.py

@@ -29,34 +29,9 @@
 memoize = t.cast(t.Callable[[T], T], functools.lru_cache(maxsize=None))
 
 
-class _WeakrefMemoized(t.Generic[T_NNF, T]):
-    def __init__(self) -> None:
-        assert t.TYPE_CHECKING, "Not a real class"
-        self.memo = NotImplemented  # type: weakref.WeakKeyDictionary[T_NNF, T]
-        self.__wrapped__ = NotImplemented  # type: t.Callable[[T_NNF], T]
-
-    @t.overload
-    def __get__(self: U, instance: None, owner: t.Type[T_NNF]) -> "U":
-        ...
-
-    @t.overload  # noqa: F811
-    def __get__(  # noqa: F811
-        self, instance: T_NNF, owner: t.Optional[t.Type[T_NNF]] = None
-    ) -> t.Callable[[], T]:
-        ...
-
-    def __get__(  # noqa: F811
-        self, instance: object, owner: object = None
-    ) -> t.Any:
-        ...
-
-    def __call__(self, sentence: T_NNF) -> T:
-        ...
-
-
 def weakref_memoize(
     func: t.Callable[[T_NNF], T]
-) -> _WeakrefMemoized[T_NNF, T]:
+) -> "_WeakrefMemoized[T_NNF, T]":
     """Make a function cache its return values using weakrefs.
 
     This makes it possible to remember sentences' properties without keeping
@@ -102,4 +77,34 @@ def wrapped(self: T_NNF) -> T:
             return ret
 
     wrapped.memo = memo  # type: ignore
+    wrapped.set = memo.__setitem__  # type: ignore
     return t.cast(_WeakrefMemoized[T_NNF, T], wrapped)
+
+
+class _WeakrefMemoized(t.Generic[T_NNF, T]):
+    """Fake class for typechecking. Should never be instantiated."""
+    def __init__(self) -> None:
+        assert t.TYPE_CHECKING, "Not a real class"
+        self.memo = NotImplemented  # type: weakref.WeakKeyDictionary[T_NNF, T]
+        self.__wrapped__ = NotImplemented  # type: t.Callable[[T_NNF], T]
+
+    @t.overload
+    def __get__(self: U, instance: None, owner: t.Type[T_NNF]) -> "U":
+        ...
+
+    @t.overload  # noqa: F811
+    def __get__(  # noqa: F811
+        self, instance: T_NNF, owner: t.Optional[t.Type[T_NNF]] = None
+    ) -> t.Callable[[], T]:
+        ...
+
+    def __get__(  # noqa: F811
+        self, instance: object, owner: object = None
+    ) -> t.Any:
+        ...
+
+    def __call__(self, sentence: T_NNF) -> T:
+        ...
+
+    def set(self, sentence: T_NNF, value: T) -> None:
+        ...

test_nnf.py

@@ -3,6 +3,7 @@
 import platform
 import shutil
 import os
+import types
 
 from pathlib import Path
 
@@ -16,6 +17,15 @@
 from nnf import (Var, And, Or, amc, dimacs, dsharp, operators,
                  tseitin, complete_models, using_kissat)
 
+memoized = [
+    method
+    for method in vars(nnf.NNF).values()
+    if isinstance(method, types.FunctionType) and hasattr(method, "memo")
+]
+assert memoized, "No memoized methods found, did the implementation change?"
+for method in memoized:
+    method.set = lambda *args: None  # type: ignore
+
 settings.register_profile('patient', deadline=2000,
                           suppress_health_check=(HealthCheck.too_slow,))
 settings.load_profile('patient')
@@ -357,7 +367,6 @@ def test_dsharp_output(fname: str):
         sentence = dsharp.load(f)
     with open(basepath + '.cnf') as f:
         clauses = dimacs.load(f)
-    nnf.NNF.decomposable.memo.clear()
     assert sentence.decomposable()
     # this is not a complete check, but clauses.models() is very expensive
     assert all(clauses.satisfied_by(model) for model in sentence.models())
@@ -464,7 +473,6 @@ def test_smoothing(sentence: nnf.NNF):
         event("Sentence not smooth yet")
         smoothed = sentence.make_smooth()
         assert type(sentence) is type(smoothed)
-        nnf.NNF.smooth.memo.clear()
         assert smoothed.smooth()
         assert sentence.equivalent(smoothed)
         assert smoothed.make_smooth() == smoothed
@@ -528,6 +536,7 @@ def test_model_counting(sentence: nnf.NNF):
 
 def test_uf20_model_counting():
     for sentence in uf20:
+        assert sentence.model_count() == len(list(sentence.models()))
         sentence.mark_deterministic()
         assert sentence.model_count() == len(list(sentence.models()))
 
@@ -546,6 +555,7 @@ def test_validity(sentence: nnf.NNF):
 
 def test_uf20_validity():
     for sentence in uf20:
+        assert not sentence.valid()
         sentence.mark_deterministic()
         assert not sentence.valid()
 
@@ -590,8 +600,6 @@ def test_is_DNF_examples():
 def test_to_MODS(sentence: nnf.NNF):
     assume(len(sentence.vars()) <= 5)
     mods = sentence.to_MODS()
-    nnf.NNF.is_MODS.memo.clear()
-    nnf.NNF._is_DNF_strict.memo.clear()
     assert mods.is_MODS()
     assert isinstance(mods, Or)
     assert mods.model_count() == len(mods.children)
@@ -807,8 +815,6 @@ def test_dsharp_compile_uf20():
         for sentence in uf20_cnf:
             compiled = dsharp.compile(sentence)
             compiled_smooth = dsharp.compile(sentence, smooth=True)
-            nnf.NNF.decomposable.memo.clear()
-            nnf.NNF.smooth.memo.clear()
             assert sentence.equivalent(compiled)
             assert sentence.equivalent(compiled_smooth)
             assert compiled.decomposable()
@@ -820,8 +826,6 @@ def test_dsharp_compile(sentence: And[Or[Var]]):
         assume(all(len(clause) > 0 for clause in sentence))
         compiled = dsharp.compile(sentence)
         compiled_smooth = dsharp.compile(sentence, smooth=True)
-        nnf.NNF.decomposable.memo.clear()
-        nnf.NNF.smooth.memo.clear()
         assert compiled.decomposable()
         assert compiled_smooth.decomposable()
         assert compiled_smooth.smooth()
@@ -869,7 +873,6 @@ def test_tseitin(sentence: nnf.NNF):
     assume(sentence.size() <= 10)
 
     T = tseitin.to_CNF(sentence)
-    nnf.NNF.is_CNF.memo.clear()
     assert T.is_CNF()
     assert T.forget_aux().equivalent(sentence)