Universal map in Standard ML

univ_map.sml

signature UNIV_MAP =
  sig
    structure Key : sig
      type 'a t
      val create : unit -> 'a t
    end

    type t

    val empty : t
    val set : t -> 'a Key.t -> 'a -> t
    val get : t -> 'a Key.t -> 'a option
  end

structure Univ_map_list :> UNIV_MAP =
  struct
    structure Key =
      struct
        type 'a t = 'a Universal.tag
        val create = Universal.tag
      end

    type t = Universal.universal list

    val empty = nil
    val set = fn l => fn tag => fn x => Universal.tagInject tag x :: l
    val get = fn l => fn tag =>
      List.foldr
        (fn (_,SOME x) => SOME x
          | (u,NONE  ) =>
              if Universal.tagIs tag u
                then SOME (Universal.tagProject tag u)
                else NONE
        )
        NONE
        l
  end

structure Univ_map_dict :> UNIV_MAP =
  struct
    structure Key =
      struct
        type 'a t = 'a Universal.tag * int

        local
          val counter = ref 0
        in
          val create = fn () =>
            (Universal.tag (), !counter before Ref.modify (fn x => x + 1) counter)
        end
      end

    structure Map =
      RedBlackMapFn (
        type ord_key = int
        val compare = Int.compare
      )

    type t = Universal.universal Map.map

    val empty = Map.empty
    val set = fn t => fn (tag, i) => fn x => Map.insert (t, i, Universal.tagInject tag x)
    val get = fn t => fn (tag, i) =>
      Option.map
        (Universal.tagProject tag)
        (Map.find (t, i))
  end