22from vyper .ast import Call , Expr
33from vyper .ir .s_expressions import parse_s_exp
44from vyper .codegen .ir_node import IRnode
5- from vyper .builtins .functions import STMT_DISPATCH_TABLE , BuiltinFunction
5+ from vyper .builtins .functions import (
6+ STMT_DISPATCH_TABLE ,
7+ DISPATCH_TABLE ,
8+ BuiltinFunction ,
9+ )
610from vyper .compiler import phases
711
812from dasy import parser
13+ from dasy .parser .utils import get_ir_type
914
1015from hy import repr
1116
1217
13- def parse_venom (expr ):
14- ir = IRnode .from_list ((parse_s_exp (repr (expr [1 ])[1 :]))[0 ])
18+ def parse_ir (expr ):
19+ # check for optional return type annotation as second element
20+ ret_type = None
21+ ir_expr = None
22+ if len (expr ) == 3 :
23+ ret_type = get_ir_type (expr [1 ].name )
24+ ir_expr = expr [2 ]
25+ elif len (expr ) == 2 :
26+ ir_expr = expr [1 ]
27+ ir = IRnode .from_list ((parse_s_exp (repr (ir_expr )[1 :]))[0 ], typ = ret_type )
1528
1629 # generate some vyper code to patch in.
1730 IDENTIFIER = f"__DASY_VENOM_BUILTIN_{ parser .next_nodeid ()} __"
@@ -21,13 +34,23 @@ def parse_venom(expr):
2134 class generated_builtin (BuiltinFunction ):
2235 _id = IDENTIFIER
2336 _inputs = ()
24- _return_type = None
37+ _return_type = ret_type
38+
39+ def fetch_call_return (self , node ):
40+ return self ._return_type
41+
42+ def infer_arg_types (self , node ):
43+ return []
2544
2645 def build_IR (self , expr , context ):
2746 return ir
2847
29- STMT_DISPATCH_TABLE [IDENTIFIER ] = generated_builtin ()
48+ if ret_type is not None :
49+ DISPATCH_TABLE [IDENTIFIER ] = generated_builtin ()
50+ gend_ast = phases .generate_ast (insert_code , 0 , "" )
51+ return gend_ast [1 ].body [0 ].value
3052
53+ STMT_DISPATCH_TABLE [IDENTIFIER ] = generated_builtin ()
3154 return phases .generate_ast (insert_code , 0 , "" )[1 ].body [0 ]
3255
3356
0 commit comments