From 9098cff929449fd8d2716b399a4d4dfe722f91f7 Mon Sep 17 00:00:00 2001
From: Marcell Vazquez-Chanlatte <mvc@linux.com>
Date: Sun, 26 Nov 2017 20:01:54 -0800
Subject: [PATCH] added var_names property

---
 stl/ast.py      | 8 +++++++-
 stl/test_ast.py | 6 ++++++
 2 files changed, 13 insertions(+), 1 deletion(-)

diff --git a/stl/ast.py b/stl/ast.py
index b64325a..f5f9858 100644
--- a/stl/ast.py
+++ b/stl/ast.py
@@ -5,7 +5,7 @@ from functools import lru_cache
 
 import funcy as fn
 import lenses
-from lenses import lens
+from lenses import lens, bind
 
 
 def flatten_binary(phi, op, dropT, shortT):
@@ -77,6 +77,12 @@ class AST(object):
     def atomic_predicates(self):
         return set(AP_lens(self).Each().collect())
 
+    @property
+    def var_names(self):
+        symbols = set(bind(self.lineqs).Each().terms.Each().collect())
+        symbols |= self.atomic_predicates
+        return set(bind(symbols).Each().id.collect())
+
     def inline_context(self, context):
         phi, phi2 = self, None
 
diff --git a/stl/test_ast.py b/stl/test_ast.py
index 90d8a2b..45eb6ad 100644
--- a/stl/test_ast.py
+++ b/stl/test_ast.py
@@ -45,3 +45,9 @@ def test_walk():
     phi = stl.parse(
         '((G[0, 1](x + y > a?)) & (F[1,2](z - x > 0))) | ((X(AP1)) U (AP2))')
     assert len(list((~phi).walk())) == 11
+
+
+def test_var_names():
+    phi = stl.parse(
+        '((G[0, 1](x + y > a?)) & (F[1,2](z - x > 0))) | ((X(AP1)) U (AP2))')
+    assert phi.var_names == {'x', 'y', 'z', 'x', 'AP1', 'AP2'}