add self_ast_simplified fold_expression

src/passes/3-self_ast_simplified/helpers.ml

@@ -1,8 +1,93 @@
 open Ast_simplified
 open Trace
 
-type mapper = expression -> expression result
+type 'a folder = 'a -> expression -> 'a result
+let rec fold_expression : 'a folder -> 'a -> expression -> 'a result = fun f init e ->
+  let self = fold_expression f in 
+  let%bind init' = f init e in
+  match e.expression with
+  | E_literal _ | E_variable _ | E_skip -> ok init'
+  | E_list lst | E_set lst | E_tuple lst | E_constant (_ , lst) -> (
+    let%bind res' = bind_fold_list self init' lst in
+    ok res'
+  )
+  | E_map lst | E_big_map lst -> (
+    let%bind res' = bind_fold_list (bind_fold_pair self) init' lst in
+    ok res'
+  )
+  | E_look_up ab | E_sequence ab | E_loop ab | E_application ab -> (
+      let%bind res' = bind_fold_pair self init' ab in
+      ok res'
+    )
+  | E_lambda { binder = _ ; input_type = _ ; output_type = _ ; result = e }
+  | E_annotation (e , _) | E_constructor (_ , e) -> (
+      let%bind res' = self init' e in
+      ok res'
+    )
+  | E_assign (_ , path , e) | E_accessor (e , path) -> (
+      let%bind res' = fold_path f init' path in
+      let%bind res' = self res' e in
+      ok res'
+    )
+  | E_matching (e , cases) -> (
+      let%bind res = self init' e in
+      let%bind res = fold_cases f res cases in
+      ok res
+    )
+  | E_record m -> (
+    let aux init'' _ expr =
+      let%bind res' = fold_expression self init'' expr in
+      ok res'
+    in
+    let%bind res = bind_fold_smap aux (ok init') m in
+    ok res
+  )
+  | E_let_in { binder = _ ; rhs ; result } -> (
+      let%bind res = self init' rhs in
+      let%bind res = self res result in
+      ok res
+    )
 
+and fold_path : 'a folder -> 'a -> access_path -> 'a result = fun f init p -> bind_fold_list (fold_access f) init p
+
+and fold_access : 'a folder -> 'a -> access -> 'a result = fun f init a ->
+  match a with
+  | Access_map e -> (
+      let%bind e' = fold_expression f init e in
+      ok e'
+    )
+  | _ -> ok init
+
+and fold_cases : 'a folder -> 'a -> matching_expr -> 'a result = fun f init m ->
+  match m with
+  | Match_bool { match_true ; match_false } -> (
+      let%bind res = fold_expression f init match_true in
+      let%bind res = fold_expression f res match_false in
+      ok res
+    )
+  | Match_list { match_nil ; match_cons = (_ , _ , cons) } -> (
+      let%bind res = fold_expression f init match_nil in
+      let%bind res = fold_expression f res cons in
+      ok res
+    )
+  | Match_option { match_none ; match_some = (_ , some) } -> (
+      let%bind res = fold_expression f init match_none in
+      let%bind res = fold_expression f res some in
+      ok res
+    )
+  | Match_tuple (_ , e) -> (
+      let%bind res = fold_expression f init e in
+      ok res
+    )
+  | Match_variant lst -> (
+      let aux init' ((_ , _) , e) =
+        let%bind res' = fold_expression f init' e in
+        ok res' in
+      let%bind res = bind_fold_list aux init lst in
+      ok res
+    )
+
+type mapper = expression -> expression result
 let rec map_expression : mapper -> expression -> expression result = fun f e ->
   let self = map_expression f in
   let%bind e' = f e in

src/passes/3-self_ast_simplified/self_ast_simplified.ml

@@ -23,3 +23,5 @@ let all_expression =
   bind_chain all_p
 
 let map_expression = Helpers.map_expression
+
+let fold_expression = Helpers.fold_expression