From 8f4e511326a8f9f57e566df09ec64676509ec2fb Mon Sep 17 00:00:00 2001
From: Marcell Vazquez-Chanlatte <mvc@linux.com>
Date: Thu, 30 Nov 2017 22:46:12 -0800
Subject: [PATCH] use lenses.Recur instead of ast_lens

---
 requirements.txt |  2 +-
 stl/ast.py       | 56 +++++++-----------------------------------------
 2 files changed, 9 insertions(+), 49 deletions(-)

diff --git a/requirements.txt b/requirements.txt
index eecadce..263db97 100644
--- a/requirements.txt
+++ b/requirements.txt
@@ -2,7 +2,7 @@
 -e git://github.com/mvcisback/multidim-threshold@master#egg=multidim-threshold
 bitarray==0.8.1
 funcy==1.9.1
-lenses==0.3.0
+lenses==0.4.0
 pandas==0.19.2
 parsimonious==0.7.0
 sympy==1.0
diff --git a/stl/ast.py b/stl/ast.py
index b368bba..4491a2f 100644
--- a/stl/ast.py
+++ b/stl/ast.py
@@ -4,7 +4,6 @@ from collections import deque, namedtuple
 from functools import lru_cache
 
 import funcy as fn
-import lenses
 from lenses import lens, bind
 
 
@@ -71,11 +70,11 @@ class AST(object):
 
     @property
     def lineqs(self):
-        return set(lineq_lens(self).Each().collect())
+        return set(lineq_lens.collect()(self))
 
     @property
     def atomic_predicates(self):
-        return set(AP_lens(self).Each().collect())
+        return set(AP_lens.collect()(self))
 
     @property
     def var_names(self):
@@ -86,11 +85,11 @@ class AST(object):
     def inline_context(self, context):
         phi, phi2 = self, None
 
-        def update(aps):
-            return tuple(context.get(ap, ap) for ap in aps)
+        def update(ap):
+            return context.get(ap, ap)
 
         while phi2 != phi:
-            phi2, phi = phi, AP_lens(phi).modify(update)
+            phi2, phi = phi, AP_lens.modify(update)(phi)
 
         return phi
 
@@ -292,47 +291,9 @@ class Param(namedtuple('Param', ['name']), AST):
         return hash(repr(self))
 
 
-def ast_lens(phi,
-             bind=True,
-             *,
-             pred=lambda _: False,
-             focus_lens=lambda _: [lens],
-             getter=False):
-    child_lenses = _ast_lens(phi, pred=pred, focus_lens=focus_lens)
-    phi = lenses.bind(phi) if bind else lens
-    return (phi.Tuple if getter else phi.Fork)(*child_lenses)
-
-
-def _ast_lens(phi, pred, focus_lens):
-    if pred(phi):
-        yield from focus_lens(phi)
-
-    if phi is None or not phi.children:
-        return
-
-    if isinstance(phi, Until):
-        child_lenses = [lens.GetAttr('arg1'), lens.GetAttr('arg2')]
-    elif isinstance(phi, NaryOpSTL):
-        child_lenses = [
-            lens.GetAttr('args')[j] for j, _ in enumerate(phi.args)
-        ]
-    else:
-        child_lenses = [lens.GetAttr('arg')]
-    for l in child_lenses:
-        yield from [l & cl for cl in _ast_lens(l.get()(phi), pred, focus_lens)]
-
-
 @lru_cache()
 def param_lens(phi, *, getter=False):
-    def focus_lens(leaf):
-        candidates = [lens.const] if isinstance(leaf, LinEq) else [
-            lens.GetAttr('interval')[0],
-            lens.GetAttr('interval')[1]
-        ]
-        return (x for x in candidates if isinstance(x.get()(leaf), Param))
-
-    return ast_lens(
-        phi, pred=type_pred(LinEq, F, G), focus_lens=focus_lens, getter=getter)
+    return bind(phi).Recur(Param)
 
 
 def type_pred(*args):
@@ -340,6 +301,5 @@ def type_pred(*args):
     return lambda x: type(x) in ast_types
 
 
-lineq_lens = fn.partial(ast_lens, pred=type_pred(LinEq), getter=True)
-AP_lens = fn.partial(ast_lens, pred=type_pred(AtomicPred), getter=True)
-and_or_lens = fn.partial(ast_lens, pred=type_pred(And, Or), getter=True)
+lineq_lens = lens.Recur(LinEq)
+AP_lens = lens.Recur(AtomicPred)