Mapping Keys

OCaml's standard library provides maps and sets, implemented using module functors. To use a map, we create a module for the type we want to use as keys via Map.Make; for a set, we create a module for the type of its elements via Set.Make. For this problem, assume we're given the following three modules.

module StringMap : Map.S with type key = string
module IntMap : Map.S with type key = int
module StringSet : Set.S with type elt = string

StringMap and IntMap are maps with string and int keys, respectively, and StringSet is a set of string values. Each supports the usual Map.S and Set.S operations, e.g., StringMap.add, StringMap.find, StringMap.bindings, StringSet.add, StringSet.elements, and so on.

Implement the function

val flip_keys_and_values : int StringMap.t -> StringSet.t IntMap.t

so that flip_keys_and_values m is the map m' gotten by making the values of m into keys. The value associated with a key v in m' is the set of keys in m that map to the value v.

let m = StringMap.of_seq (List.to_seq [("a", 1); ("b", 2); ("c", 1); ("d", 3)])
let out = flip_keys_and_values m
let expected =
  IntMap.of_seq (List.to_seq
    [(1, StringSet.of_list ["a"; "c"]);
     (2, StringSet.of_list ["b"]);
     (3, StringSet.of_list ["d"])])
let _ = assert (IntMap.equal StringSet.equal out expected)

let _ = assert
    (IntMap.equal StringSet.equal
       (flip_keys_and_values StringMap.empty) IntMap.empty)