--- a/pysmcl/ideal_functionality.py	Wed Dec 23 14:35:43 2009 +0100
+++ b/pysmcl/ideal_functionality.py	Wed Dec 23 14:40:26 2009 +0100
@@ -1,5 +1,7 @@
 import inspect
-from twisted.internet.defer import inlineCallbacks
+from twisted.internet.defer import _DefGen_Return
+from pysmcl.defer_wrapper import inlineCallbacks
+from twisted.internet.defer import returnValue
 
 import pysmcl.ast_wrapper as ast
 import pysmcl.compatibility_check
@@ -11,12 +13,21 @@
 
 debug = True
 
+def catcher(f):
+    def c(*args, **kwargs):
+        try:
+            return f(*args, **kwargs)
+        except _DefGen_Return, e:
+            return e.value
+    return c
 
 def ideal_functionality(range={}, secrets=[]):
     def ideal_functionality_helper(f):
         source_str = inspect.getsource(f)
         source_ast = ast.parse(source_str) # Returns a module
         function_ast = source_ast.body[0] # We want only the function
+        # for i in ast.walk(function_ast):
+        #     print i, getattr(i, "lineno", None)
         if(debug):
             pysmcl.pretty_print.pprint(source_ast)
             print "-"*80
@@ -31,6 +42,7 @@
         bad_calls(function_ast)
 
         pysmcl.secret_annotator.non_compromising_functions.add(function_ast.name)
+        pysmcl.secret_annotator.annotated_functions.add(function_ast.name)
         if(pysmcl.secret_annotator.returns_secret(function_ast)):
             pysmcl.secret_annotator.secret_functions.add(function_ast.name)
 
@@ -40,8 +52,13 @@
         if(debug):
             pysmcl.pretty_print.pprint(source_ast)
         g = f.__globals__
+        g["returnValue"] = returnValue
         env = {}
-        ast.increment_lineno(t, 3)
+        ast.increment_lineno(t, 10)
+        # for i in ast.walk(function_ast):
+        #     print i, getattr(i, "lineno", None)
+        #     i.lineno = 3
+        # print ast.dump(function_ast)
         exec compile(ast.Module([t]), f.__code__.co_filename, "exec") in g, env
         f = env[function_ast.name]
         return inlineCallbacks(f)