From 7776c837d6f4565520edb2dd70211d52406de471 Mon Sep 17 00:00:00 2001
From: Max Nordlund gmail <max.nordlund@gmail.com>
Date: Tue, 25 Oct 2022 22:29:14 +0200
Subject: [PATCH] WIP: map support

---
 include/proper.hrl                  |  6 ++--
 src/proper_types.erl                | 43 ++++++++++++++++++-----
 src/proper_typeserver.erl           | 53 +++++++++++++++++++++++++++++
 test/proper_exported_types_test.erl |  3 --
 test/proper_tests.erl               |  6 +++-
 5 files changed, 95 insertions(+), 16 deletions(-)

diff --git a/include/proper.hrl b/include/proper.hrl
index 00219779..996377a0 100644
--- a/include/proper.hrl
+++ b/include/proper.hrl
@@ -47,9 +47,9 @@
 %%------------------------------------------------------------------------------
 
 -import(proper_types, [integer/2, float/2, atom/0, binary/0, binary/1,
-		       bitstring/0, bitstring/1, list/1, vector/2, union/1,
-		       weighted_union/1, tuple/1, loose_tuple/1, exactly/1,
-		       fixed_list/1, function/2, map/2, any/0]).
+		       bitstring/0, bitstring/1, list/1, map/1, map/2, map_union/2,
+			   vector/2, union/1, weighted_union/1, tuple/1, loose_tuple/1,
+			   exactly/1, fixed_list/1, fixed_map/1, function/2, any/0]).
 
 
 %%------------------------------------------------------------------------------
diff --git a/src/proper_types.erl b/src/proper_types.erl
index c663b63a..3b79ae57 100644
--- a/src/proper_types.erl
+++ b/src/proper_types.erl
@@ -142,13 +142,13 @@
 -export([integer/2, float/2, atom/0, binary/0, binary/1, bitstring/0,
 	 bitstring/1, list/1, vector/2, union/1, weighted_union/1, tuple/1,
 	 loose_tuple/1, exactly/1, fixed_list/1, fixed_map/1, function/2, map/0,
-     map/2, any/0, shrink_list/1, safe_union/1, safe_weighted_union/1]).
+     map/1, map/2, any/0, shrink_list/1, safe_union/1, safe_weighted_union/1]).
 -export([integer/0, non_neg_integer/0, pos_integer/0, neg_integer/0, range/2,
 	 float/0, non_neg_float/0, number/0, boolean/0, byte/0, char/0, nil/0,
 	 list/0, tuple/0, string/0, wunion/1, term/0, timeout/0, arity/0]).
 -export([int/0, nat/0, largeint/0, real/0, bool/0, choose/2, elements/1,
 	 oneof/1, frequency/1, return/1, default/2, orderedlist/1, function0/1,
-	 function1/1, function2/1, function3/1, function4/1,
+	 function1/1, function2/1, function3/1, function4/1, map_union/1,
 	 weighted_default/2]).
 -export([resize/2, non_empty/1, noshrink/1]).
 
@@ -1120,28 +1120,49 @@ function_is_instance(Type, X) ->
 map() ->
     ?LAZY(map(any(), any())).
 
+%% @doc A map whose keys and values are defined by the given `Map'.
+%%
+%% Shrinks towards the empty map. That is, all keys are assumed to be optional.
+%%
+%% Also written simply as a {@link maps. map}.
+-spec map(#{Key::raw_type() => Value::raw_type()}) -> proper_types:type().
+map(Map) when is_map(Map) ->
+    MapType = maps:map(fun(_Key, Value) -> cook_outer(Value) end, Map),
+    ?CONTAINER([
+        {generator, {typed, fun map_gen/1}},
+        {is_instance, {typed, fun map_is_instance/2}},
+        {internal_types, MapType},
+        {get_length, fun maps:size/1},
+        {join, fun maps:merge/2},
+        {get_indices, fun fixed_map_get_keys/2},
+        {remove, fun maps:remove/2},
+        {retrieve, fun maps:get/2},
+        {update, fun maps:update/3}
+    ]).
+
 %% @doc A map whose keys are defined by the generator `K' and values
 %% by the generator `V'.
 -spec map(K::raw_type(), V::raw_type()) -> proper_types:type().
 map(K, V) ->
     ?LET(L, list({K, V}), maps:from_list(L)).
 
+%% @doc A map merged from the given map generators.
+-spec map_union([Map::raw_type()]) -> proper_types:type().
+map_union(RawMaps) when is_list(RawMaps) ->
+    ?LET(Maps, RawMaps, lists:foldl(fun maps:merge/2, #{}, Maps)).
+
 %% @doc A map whose keys and values are defined by the given `Map'.
 %% Also written simply as a {@link maps. map}.
 -spec fixed_map(#{Key::raw_type() => Value::raw_type()}) -> proper_types:type().
-% fixed_map(Map) when is_map(Map) ->
-%     Pairs = maps:to_list(Map),
-%     ?LET(L, fixed_list(Pairs), maps:from_list(L)).
-
 fixed_map(Map) when is_map(Map) ->
+    MapType = maps:map(fun(_Key, Value) -> cook_outer(Value) end, Map),
     ?CONTAINER([
         {generator, {typed, fun map_gen/1}},
         {is_instance, {typed, fun map_is_instance/2}},
-        {internal_types, Map},
+        {internal_types, MapType},
         {get_length, fun maps:size/1},
         {join, fun maps:merge/2},
-        {get_indices, fun maps:keys/1},
-        {remove, fun maps:remove/2},
+        {get_indices, fun fixed_map_get_keys/2},
         {retrieve, fun maps:get/2},
         {update, fun maps:update/3}
     ]).
@@ -1184,6 +1205,10 @@ map_all_internal(Fun, none, Result) when is_function(Fun, 2) andalso is_boolean(
 map_all_internal(Fun, {Key, Value, NextIterator}, true) when is_function(Fun, 2) ->
     map_all_internal(Fun, NextIterator, Fun(Key, Value)).
 
+fixed_map_get_keys(Type, _X) ->
+    Map = get_prop(internal_types, Type),
+    maps:keys(Map).
+
 %% @doc All Erlang terms (that PropEr can produce). For reasons of efficiency,
 %% functions are never produced as instances of this type.<br />
 %% CAUTION: Instances of this type are expensive to produce, shrink and instance-
diff --git a/src/proper_typeserver.erl b/src/proper_typeserver.erl
index c5ea8112..120f6936 100644
--- a/src/proper_typeserver.erl
+++ b/src/proper_typeserver.erl
@@ -1648,6 +1648,8 @@ convert(_Mod, {type,_,nonempty_string,[]}, State, _Stack, _VarDict) ->
     {ok, {simple,proper_types:non_empty(proper_types:string())}, State};
 convert(_Mod, {type,_,map,any}, State, _Stack, _VarDict) ->
     {ok, {simple,proper_types:map()}, State};
+convert(Mod, {type,_,map,Fields}, State, Stack, VarDict) ->
+	convert_map(Mod, Fields, State, Stack, VarDict);
 convert(_Mod, {type,_,tuple,any}, State, _Stack, _VarDict) ->
     {ok, {simple,proper_types:tuple()}, State};
 convert(Mod, {type,_,tuple,ElemForms}, State, Stack, VarDict) ->
@@ -1787,6 +1789,57 @@ convert_normal_rec_list(RecFun, RecArgs, NonEmpty) ->
     NewRecArgs = clean_rec_args(RecArgs),
     {NewRecFun, NewRecArgs}.
 
+-spec convert_map(mod_name(), [Field], state(), stack(), var_dict()) ->
+	rich_result2(ret_type(), state())
+when
+	Field :: {type, erl_anno:anno(), map_field_assoc, [abs_type()]}
+		| {type, erl_anno:anno(), map_field_exact, [abs_type()]}.
+convert_map(Mod, Fields, State1, Stack, VarDict) ->
+	{AbstractRequiredFields, AbstractOptionalFields} = lists:partition(
+		fun ({type, _, map_field_exact, _FieldType}) ->
+				true;
+			({type, _, map_field_assoc, _FieldType}) ->
+				false
+		end,
+		Fields
+	),
+	case process_map_fields(Mod, AbstractRequiredFields, State1, Stack, VarDict) of
+		{ok, RequiredFields, State2} ->
+			case process_map_fields(Mod, AbstractOptionalFields, State2, Stack, VarDict) of
+				{ok, OptionalFields, State3} ->
+					Required = proper_types:fixed_map(maps:from_list(RequiredFields)),
+					Optional = proper_types:map(maps:from_list(OptionalFields)),
+					{ok, {simple, proper_types:map_union([Required, Optional])}, State3};
+				{error, Reason} ->
+					{error, Reason}
+			end;
+		{error, Reason} ->
+			{error, Reason}
+	end.
+
+process_map_fields(Mod, AbstractFields, State, Stack, VarDict) ->
+	Process =
+		fun ({type, _, _, RawFieldTypes}, {ok, Fields, State1}) when
+				length(RawFieldTypes) =:= 2
+			->
+				case process_list(
+					Mod, RawFieldTypes, State1, [map | Stack], VarDict
+				) of
+					{ok, FieldTypes, State2} ->
+						{ok, [list_to_tuple(FieldTypes) | Fields], State2};
+					{error, Reason} ->
+						{error, Reason}
+				end;
+			(_FieldTypes, {error, Reason}) ->
+				{error, Reason}
+	end,
+	case lists:foldl(Process, {ok, [], State}, AbstractFields) of
+		{ok, ReverseFields, NewState} ->
+			{ok, lists:reverse(ReverseFields), NewState};
+		{error, Reason} ->
+			{error, Reason}
+	end.
+
 -spec convert_tuple(mod_name(), [abs_type()], boolean(), state(), stack(),
 		    var_dict()) -> rich_result2(ret_type(),state()).
 convert_tuple(Mod, ElemForms, ToList, State, Stack, VarDict) ->
diff --git a/test/proper_exported_types_test.erl b/test/proper_exported_types_test.erl
index 7b7e801b..e90bc27f 100644
--- a/test/proper_exported_types_test.erl
+++ b/test/proper_exported_types_test.erl
@@ -40,9 +40,6 @@
 %%
 %% Still, the test is currently not 100% there.
 %% TODOs:
-%%   - Eliminate the 12 errors that `proper_typeserver:demo_translate_type/2`
-%%     currently returns. (Three of these errors are due to the incomplete
-%%     handling of maps.)
 %%   - Handle symbolic instances (the {'$call', ...} case below).
 %%
 
diff --git a/test/proper_tests.erl b/test/proper_tests.erl
index 6960872c..9baea8a0 100644
--- a/test/proper_tests.erl
+++ b/test/proper_tests.erl
@@ -410,6 +410,10 @@ simple_types_with_data() ->
      {[], [[]], [], [[a],[1,2,3]], "[]"},
      {fixed_list([neg_integer(),pos_integer()]), [[-12,32],[-1,1]], [-1,1],
       [[0,0]], none},
+     {map(#{key => value, pos_integer() => neg_integer()}),
+      [#{key => value, 1 => -1}], #{}, [not_a_map], none},
+     {fixed_map(#{key => value, pos_integer() => neg_integer()}),
+      [#{key => value, 3 => -3}], #{key => value, 1 => -1}, [not_a_map], none},
      {[atom(),integer(),atom(),float()], [[forty_two,42,forty_two,42.0]],
       ['',0,'',0.0], [[proper,is,licensed],[under,the,gpl]], none},
      {[42 | list(integer())], [[42],[42,44,22]], [42], [[],[11,12]], none},
@@ -773,7 +777,7 @@ cant_generate_test_() ->
     [?_test(assert_cant_generate(Type)) || Type <- impossible_types()].
 
 proper_exported_types_test_() ->
-    [?_assertEqual({[],12}, proper_exported_types_test:not_handled())].
+    [?_assertEqual({[],0}, proper_exported_types_test:not_handled())].
 
 %%------------------------------------------------------------------------------
 %% Verify that failing constraints are correctly reported