--- a/pysmcl/editor_info.py	Wed Dec 16 17:43:49 2009 +0100
+++ b/pysmcl/editor_info.py	Wed Dec 16 18:18:01 2009 +0100
@@ -5,7 +5,8 @@
 import pysmcl.pretty_print as pp
 from pysmcl.range_analysis import RangeAnalysis
 from pysmcl.secret_annotator import secret_analysis, expr_secret
-import pysmcl.secret_ifs
+import pysmcl.secret_ifs as secret_ifs
+
 
 def main():
     try:
@@ -18,46 +19,64 @@
             if decorator is None:
                 continue
             secret_analysis(i)
-            
+
             RangeAnalysis(pysmcl.setup.Zp.modulus).apply(i, {})
 
-            transform_ifs = pysmcl.secret_ifs.TransformIfs()
-            t = transform_ifs.visit(i)
+            transform_ifs = secret_ifs.TransformIfs()
+            i = transform_ifs.visit(i)
             while transform_ifs.changed:
                 secret_analysis(i)
                 RangeAnalysis(pysmcl.setup.Zp.modulus).apply(i, {})
                 transform_ifs.reset()
-                t = transform_ifs.visit(i)
+                i = transform_ifs.visit(i)
             secret_analysis(i)
             func_string = pp.pprint_string(i)
             print "!transformed: %d %d %d (%s)" % (i.lineno, i.col_offset,
-                                                len(i.name), func_string) 
+                                                len(i.name), func_string)
+
+    class ValueVisitor(ast.NodeVisitor):
+
+        def __init__(self):
+            self.inside_function_def = False
+            pass
 
-    # Graph(Flow().to_dot(prog.body[1])).to_dot()
-    for top_level_statement in prog.body:
-        if(isinstance(top_level_statement, ast.FunctionDef)):
-            decorator = ast.get_ideal_functionality(top_level_statement)
+        def visit_FunctionDef(self, node):
+            decorator = ast.get_ideal_functionality(node)
             if decorator is None:
-                continue
+                return
+            self.inside_function_def = True
+            for i in node.body:
+                self.visit(i)
+            self.inside_function_def = False
+
+        def visit_Assign(self, node):
+            if(secret_ifs.is_synthetic(node)):
+                return
+            self.generic_visit(node)
 
-            for statement in top_level_statement.body:
-                if getattr(statement, "is_syntetic", False):
-                    continue
-                for j in ast.walk(statement):
-                    if isinstance(j, ast.stmt):
-                        e = j.out_values["secret"]
-                        r = j.out_values["range"]
-                    if isinstance(j, ast.expr):
-                        if expr_secret(j, e):
-                            s = pp.expr_string(j)
-                            print "!secret: %d %d %d (%s)" % (j.lineno, j.col_offset,
-                                                len(s), s)
-                    if(isinstance(j, ast.Name) and not 
-                       (isinstance(j.parent, ast.Call) and j is j.parent.func)):
+        def generic_visit(self, node):
+            if self.inside_function_def:
+                if isinstance(node, ast.stmt):
+                    self.e = node.out_values["secret"]
+                    self.r = node.out_values["range"]
+                if(isinstance(node, ast.Name)
+                   and not (isinstance(node.parent, ast.Call)
+                            and node is node.parent.func)):
+                    print("!values: %d %d %d (%s)"
+                          % (node.lineno, node.col_offset,
+                             len(pp.expr_string(node)),
+                             node.id + " is in " +str(self.r[node.id])))
+                if(isinstance(node, ast.expr)
+                   and not isinstance(getattr(node, 'ctx', None), ast.Store)):
+                    if expr_secret(node, self.e):
+                        s = pp.expr_string(node)
+                        print "!secret: %d %d %d (%s)" % (node.lineno,
+                                                          node.col_offset,
+                                                          len(s), s)
+                        return # We do not want all subexpressions to be marked
+            ast.NodeVisitor.generic_visit(self, node)
 
-                        print "!values: %d %d %d (%s)" % (j.lineno, j.col_offset,
-                                                len(pp.expr_string(j)), j.id + " is in " +str(r[j.id]))
-
+    ValueVisitor().visit(prog)
 
 if __name__ == "__main__":
     main()