--- a/pysmcl/flow.py	Tue May 26 11:35:52 2009 +0200
+++ b/pysmcl/flow.py	Tue May 26 11:36:41 2009 +0200
@@ -6,26 +6,37 @@
 import pysmcl.secret_ifs
 
 
-def analyze(node, join, combine, key):
+def analyze(function, join, combine, key, bottom):
     """
-    Doing a simple forward analysis using the iterative worklist
+    Doing a simple intra procedural forward analysis using the iterative worklist
     algorithm. Parametrized by the functions join and combine, and
     with the value initial.  The lattice-points are stored in each
     statement-node in_values and out_values, under the *key*.
+
+    function (FunctionDef) The function to analyze
+    join (function) A function which takes a set of nodes to be joined
+    combine (function) A function which implement the least upper bound or
+        the greates lower bound
+    key (String) A string used to index into the out_values dictionary
+    bottom (function) A function which computes the bottom element of the lattice
     """
     # initialization
     worklist = []
 
-    Flow().flow(node)
-    for child in ast.walk(node):
-        if(isinstance(child, ast.stmt)):
-            child.in_values[key] = combine(child, set())
-            worklist.append(child)
+    Flow().flow(function)
+#     function.in_values[key] = combine(function, bottom())
+    worklist.append(function)
+#     for child in ast.walk(function):
+#         if(isinstance(child, ast.stmt)):
+#             child.in_values[key] = combine(child, bottom())
+#             print "Appending", child
+#             worklist.append(child)
 
     # main cycle
     while len(worklist) != 0:
         x = worklist.pop()
-        oldout = x.out_values[key]
+#        print "Working", x, x.parents
+        oldout = x.out_values[key]        
         x.out_values[key] = combine(x, join(x.parents))
         if oldout != x.out_values[key]:
             for child in x.children:
@@ -41,6 +52,9 @@
 
     def flow(self, function):
         #function.body.append(ast.Pass)
+        function.parents = set()
+        function.children = [function.body[0]]
+        function.body[0].parents = [function]
         self.flow_body(function.body)
         # Todo - edges from function returns...
         # Neccesary?

--- a/pysmcl/secret_annotator.py	Tue May 26 11:35:52 2009 +0200
+++ b/pysmcl/secret_annotator.py	Tue May 26 11:36:41 2009 +0200
@@ -77,10 +77,12 @@
        node.value.func.id == "mark_secret"):
         # Todo - handle non-names...
         ins = ins | set([i.id for i in node.value.args])
+    if (isinstance(node, ast.FunctionDef)):
+        ins = ins | set([arg.id for arg in node.args.args])
     return ins
 
 
 def secret_analysis(function):
     arguments = set(i.id for i in function.args.args)
     function.body[0].imported_secrets = arguments
-    pysmcl.flow.analyze(function, secret_join, secret_combine, "secret")
+    pysmcl.flow.analyze(function, secret_join, secret_combine, "secret", lambda: set())

--- a/pysmcl/secret_ifs.py	Tue May 26 11:35:52 2009 +0200
+++ b/pysmcl/secret_ifs.py	Tue May 26 11:36:41 2009 +0200
@@ -42,23 +42,6 @@
 
 
 class TransformIfs(ast.NodeTransformer):
-    """
-    >>> from ast import *
-    >>> from pysmcl.pretty_print import *
-    >>> from pysmcl.ideal_functionality import init_statements
-    >>> from pysmcl.secret_ifs import TransformIfs
-    >>> from pysmcl.secret_annotator import secret_analysis
-    >>> prog = parse("def f(x):\n\tif(x):\n\t\ta=1\n\telse:\n\t\ta=2")
-    >>> trans = TransformIfs()
-    >>> init_statements(prog)
-    >>> secret_analysis(prog.body[0])
-    >>> prog = trans.visit(prog)
-    >>> pprint(prog)
-    def f(x):
-        cond0 = x
-        a = cond0 * 1 + (1 - cond0) * 2
-
-    """
     cond_counter = 0
 
     def __init__(self):
@@ -107,3 +90,6 @@
             return r # A list of statements will be merged into the list
         else:
             return node
+
+        
+