From 7aed260d5d2b112cf86136114608711d95d4777b Mon Sep 17 00:00:00 2001
From: Marcell Vazquez-Chanlatte <mvc@linux.com>
Date: Tue, 14 Nov 2017 09:58:56 -0800
Subject: [PATCH] fixed inline context + added tests for inline context

---
 stl/ast.py        | 11 +++++++++++
 stl/test_utils.py | 26 ++++++++++++++++++++++++++
 stl/utils.py      | 14 +-------------
 3 files changed, 38 insertions(+), 13 deletions(-)

diff --git a/stl/ast.py b/stl/ast.py
index e16a934..eddf9ed 100644
--- a/stl/ast.py
+++ b/stl/ast.py
@@ -78,6 +78,17 @@ class AST(object):
     def atomic_predicates(self):
         return set(AP_lens(self).Each().collect())
 
+    def inline_context(self, context):
+        phi, phi2 = self, None
+
+        def update(aps):
+            return tuple(context.get(ap, ap) for ap in aps)
+
+        while phi2 != phi:
+            phi2, phi = phi, AP_lens(phi).modify(update)
+        # TODO: this is hack to flatten the AST. Fix!
+        return phi
+
 
 class _Top(AST):
     __slots__ = ()
diff --git a/stl/test_utils.py b/stl/test_utils.py
index 250ad54..465a93e 100644
--- a/stl/test_utils.py
+++ b/stl/test_utils.py
@@ -4,6 +4,16 @@ from stl.hypothesis import SignalTemporalLogicStrategy
 from hypothesis import given
 
 
+CONTEXT = {
+    stl.parse('AP1'): stl.parse('F(x > 4)'),
+    stl.parse('AP2'): stl.parse('(AP1) U (AP1)'),
+    stl.parse('AP3'): stl.parse('y < 4'),
+    stl.parse('AP4'): stl.parse('y < 3'),
+    stl.parse('AP5'): stl.parse('y + x > 4'),
+}
+APS = set(CONTEXT.keys())
+
+
 @given(SignalTemporalLogicStrategy)
 def test_f_neg_or_canonical_form(phi):
     phi2 = stl.utils.f_neg_or_canonical_form(phi)
@@ -11,3 +21,19 @@ def test_f_neg_or_canonical_form(phi):
     assert phi2 == phi3
     assert not any(
         isinstance(x, (stl.ast.G, stl.ast.And)) for x in phi2.walk())
+
+
+def test_inline_context_rigid():
+    phi = stl.parse('G(AP1)')
+    phi2 = phi.inline_context(CONTEXT)
+    assert phi2 == stl.parse('G(F(x > 4))')
+
+    phi = stl.parse('G(AP2)')
+    phi2 = phi.inline_context(CONTEXT)
+    assert phi2 == stl.parse('G((F(x > 4)) U (F(x > 4)))')
+
+
+@given(SignalTemporalLogicStrategy)
+def test_inline_context(phi):
+    phi2 = phi.inline_context(CONTEXT)
+    assert not (APS & phi2.atomic_predicates)
diff --git a/stl/utils.py b/stl/utils.py
index 6e46efd..c46e3f1 100644
--- a/stl/utils.py
+++ b/stl/utils.py
@@ -6,7 +6,7 @@ from lenses import bind
 
 import stl.ast
 from stl.ast import (And, F, G, Interval, LinEq, Neg,
-                     Or, AP_lens, Next, Until, AtomicPred,
+                     Or, Next, Until, AtomicPred,
                      _Top, _Bot)
 from stl.types import STL
 
@@ -50,18 +50,6 @@ def linear_stl_lipschitz(phi):
     return float(max(map(_lineq_lipschitz, phi.lineqs)))
 
 
-def inline_context(phi, context):
-    phi2 = None
-
-    def update(ap):
-        return context.get(ap, ap)
-
-    while phi2 != phi:
-        phi2, phi = phi, AP_lens(phi).modify(update)
-    # TODO: this is hack to flatten the AST. Fix!
-    return stl.parse(str(phi))
-
-
 op_lookup = {
     ">": op.gt,
     ">=": op.ge,