--- a/pysmcl/runtime_sugar.py	Wed Dec 23 14:42:18 2009 +0100
+++ b/pysmcl/runtime_sugar.py	Wed Dec 23 14:43:19 2009 +0100
@@ -1,4 +1,5 @@
 import pysmcl.ast_wrapper as ast
+import pysmcl.secret_annotator
 
 
 """
@@ -22,7 +23,6 @@
 
     name = make_runtime_name(f, ctx=ast.Param())
     f.args.args.insert(0, name)
-    print f.args
 
     class Rewriter(ast.NodeTransformer):
 
@@ -35,7 +35,9 @@
                 ast.copy_location(r, node)
                 return r
             return node
+
         def visit_Call(self, node):
+            self.generic_visit(node)
             runtime = make_runtime_name(node.func)
             if(self.is_function(node, "output")):
                 r = ast.Attribute(value=runtime,
@@ -52,21 +54,36 @@
                                   ctx=ast.Load())
                 ast.copy_location(r, node.func)
                 node.func = r
-            if(self.is_function(node, "get")):
+                return node
+            if(node.func.id in pysmcl.secret_annotator.annotated_functions):
                 node.args.insert(0, runtime)
+                num = ast.copy_location(ast.Num(0), node)
+                slice = ast.copy_location(ast.Index(num), node)
+#                return ast.copy_location(ast.Subscript(value=node,
+#                                                       slice=slice,
+#                                                       ctx=ast.Load()),node)
             return node
 
         def visit_Return(self, node):
-            returnValue = ast.Name(id="returnValue", ctx=ast.Load)
-            e = ast.Call(func=ast.copy_location(returnValue, node),
-                         args=[node.value],
-                         keywords=[],
-                         starargs=[],
-                         kwargs=[])
-            r = ast.Expr(value=ast.copy_location(e, node))
-            return ast.copy_location(r, node)
+            returnValue = ast.copy_location(
+                ast.Name(id="returnValue", ctx=ast.Load()), node)
+#            lst = ast.copy_location(ast.List(elts=[node.value],
+#                                             ctx=ast.Load()), node)
+            call = ast.copy_location(ast.Call(func=returnValue,
+                                              args=[node.value],
+                                              keywords=[],
+                                              starargs=None,
+                                              kwargs=None), node)
+            
+            r = ast.copy_location(ast.Expr(value=call), node)
+            return r
 
         def is_function(self, node, fname):
             return isinstance(node.func, ast.Name) and node.func.id == fname
 
     Rewriter().visit(f)
+    y = ast.copy_location(ast.Yield(value=ast.copy_location(ast.Num(0), f)), f)
+    e = ast.copy_location(ast.Expr(value=y), f)
+    if_stm = ast.copy_location(ast.If(test=ast.copy_location(ast.Num(1), f), body=f.body,orelse=[e]), f)
+    f.body = [if_stm]
+    ast.fix_ast_parents(f)