From d7a16c47c1361377257121e2ac872dadcd4400ef Mon Sep 17 00:00:00 2001
From: galfour <gabriel.alfour@gmail.com>
Date: Sat, 20 Jul 2019 16:42:34 +0200
Subject: [PATCH] add iterators for maps

---
 src/contracts/map.ligo        | 11 +++++++++++
 src/operators/operators.ml    | 20 ++++++++------------
 src/test/integration_tests.ml | 10 ++++++++++
 3 files changed, 29 insertions(+), 12 deletions(-)

diff --git a/src/contracts/map.ligo b/src/contracts/map.ligo
index bcc2a8005..f0576bf54 100644
--- a/src/contracts/map.ligo
+++ b/src/contracts/map.ligo
@@ -31,3 +31,14 @@ const bm : foobar = map
   120 -> 23 ;
   421 -> 23 ;
 end
+
+function iter_op (const m : foobar) : int is
+  var r : int := 0 ;
+  function aggregate (const i : int ; const j : int) : unit is block { r := r + i + j } with unit ;
+  block {
+    map_iter(m , aggregate) ;
+  } with r ;
+
+function map_op (const m : foobar) : foobar is
+  function increment (const i : int ; const j : int) : int is block { skip } with j + 1 ;
+  block { skip } with map_map(m , increment) ;
diff --git a/src/operators/operators.ml b/src/operators/operators.ml
index 6a30913f5..67c0cdb28 100644
--- a/src/operators/operators.ml
+++ b/src/operators/operators.ml
@@ -256,22 +256,18 @@ module Typer = struct
     let%bind () = assert_type_value_eq (src, k) in
     ok @@ t_option dst ()
 
-  let map_iter : typer = typer_2 "MAP_ITER" @@ fun f m ->
+  let map_iter : typer = typer_2 "MAP_ITER" @@ fun m f ->
     let%bind (k, v) = get_t_map m in
-    let%bind (arg_1 , res) = get_t_function f in
-    let%bind (arg_2 , res') = get_t_function res in
-    let%bind () = assert_eq_1 arg_1 k in
-    let%bind () = assert_eq_1 arg_2 v in
-    let%bind () = assert_eq_1 res' (t_unit ()) in
+    let%bind (arg , res) = get_t_function f in
+    let%bind () = assert_eq_1 arg (t_pair k v ()) in
+    let%bind () = assert_eq_1 res (t_unit ()) in
     ok @@ t_unit ()
 
-  let map_map : typer = typer_2 "MAP_MAP" @@ fun f m ->
+  let map_map : typer = typer_2 "MAP_MAP" @@ fun m f ->
     let%bind (k, v) = get_t_map m in
-    let%bind (arg_1 , res) = get_t_function f in
-    let%bind (arg_2 , res') = get_t_function res in
-    let%bind () = assert_eq_1 arg_1 k in
-    let%bind () = assert_eq_1 arg_2 v in
-    ok @@ res'
+    let%bind (arg , res) = get_t_function f in
+    let%bind () = assert_eq_1 arg (t_pair k v ()) in
+    ok @@ t_map k res ()
 
   let map_fold : typer = typer_2 "MAP_FOLD" @@ fun f m ->
     let%bind (k, v) = get_t_map m in
diff --git a/src/test/integration_tests.ml b/src/test/integration_tests.ml
index d3f54421e..0a978b6e5 100644
--- a/src/test/integration_tests.ml
+++ b/src/test/integration_tests.ml
@@ -348,6 +348,16 @@ let map () : unit result =
     let expected = ez [23, 23] in
     expect_eq program "rm" input expected
   in
+  let%bind () =
+    let input = ez [(1 , 10) ; (2 , 20) ; (3 , 30) ] in
+    let expected = e_int 66 in
+    expect_eq program "iter_op" input expected
+  in
+  let%bind () =
+    let input = ez [(1 , 10) ; (2 , 20) ; (3 , 30) ] in
+    let expected = ez [(1 , 11) ; (2 , 21) ; (3 , 31) ] in
+    expect_eq program "map_op" input expected
+  in
   ok ()
 
 let list () : unit result =