From 268d8bdd8ff4e94756cf186f30e632c12349a206 Mon Sep 17 00:00:00 2001
From: Marcell Vazquez-Chanlatte <mvc@linux.com>
Date: Wed, 22 Nov 2017 13:03:32 -0800
Subject: [PATCH] add ability to distribute next (push to bottom of AST)

---
 stl/test_utils.py | 10 ++++++++++
 stl/utils.py      | 19 +++++++++++++++++--
 2 files changed, 27 insertions(+), 2 deletions(-)

diff --git a/stl/test_utils.py b/stl/test_utils.py
index ae7134b..7aaf269 100644
--- a/stl/test_utils.py
+++ b/stl/test_utils.py
@@ -86,6 +86,16 @@ def test_discretize():
     phi = stl.parse('G[0.3, 1.2]((AP1) U (AP2))')
     assert not stl.utils.is_discretizable(phi, dt)
 
+    phi = stl.parse('G[0.3, 0.6](~(F[0, 0.3](A)))')
+    assert stl.utils.is_discretizable(phi, dt)
+    phi2 = stl.utils.discretize(phi, dt, distribute=True)
+    phi3 = stl.utils.discretize(phi2, dt, distribute=True)
+    assert phi2 == phi3
+    assert phi2 == stl.parse(
+        '(~((X(A)) ∨ (X(X(A))))) ∧ (~((X(X(A))) ∨ (X(X(X(A))))))')
+
+
+
 
 def test_scope():
     dt = 0.3
diff --git a/stl/utils.py b/stl/utils.py
index a9171a6..341fe56 100644
--- a/stl/utils.py
+++ b/stl/utils.py
@@ -122,8 +122,9 @@ def scope(phi, dt, *, _t=0):
 
 
 @require_discretizable
-def discretize(phi, dt):
-    return _discretize(phi, dt)
+def discretize(phi, dt, distribute=False):
+    phi = _discretize(phi, dt)
+    return _distribute_next(phi) if distribute else phi
 
 
 def _discretize(phi, dt):
@@ -151,6 +152,20 @@ def _interval_discretizable(itvl, dt):
     return np.isclose(l, round(l)) and np.isclose(u, round(u))
 
 
+def _distribute_next(phi, i=0):
+    if isinstance(phi, (LinEq, AtomicPred)):
+        return stl.utils.next(phi, i=i)
+    elif isinstance(phi, Next):
+        return _distribute_next(phi.arg, i=i+1)
+
+    children = tuple(_distribute_next(c, i) for c in phi.children)
+
+    if isinstance(phi, (And, Or)):
+        return bind(phi).args.set(children)
+    elif isinstance(phi, (Neg, Next)):
+        return bind(phi).arg.set(children[0])
+
+
 def is_discretizable(phi, dt):
     if any(c for c in phi.walk() if isinstance(c, Until)):
         return False