Persistent Map implemented as a Concestor data structure

ConcestorMap.fs

// This code implements a persistent dictionary similar to Map except the
// internal representation is a Dictionary + diffs. This offers different
// performance tradeoffs and, in particular, much faster add and remove when
// used linearly.
//
// For more details, see the following articles from the F# Journal:
//
// http://fsharpnews.blogspot.co.uk/2017/10/the-concestor-set.html
// http://fsharpnews.blogspot.co.uk/2017/10/a-simple-concestor-dictionary.html

open System.Collections.Generic

type Operation =
  | Add
  | Remove
  | Replace

type Ancestry<'k, 'v> =
  | Delta of Operation * 'k * 'v * 'v * Concestor<'k, 'v>
  | HashTable of Dictionary<'k, 'v>
and Concestor<'k, 'v> =
  { mutable Ancestry: Ancestry<'k, 'v>
    LockObj: obj }

let create s = { Ancestry = HashTable s; LockObj = obj() }

let empty<'k, 'v when 'k: equality> =
  create(Dictionary<'k, 'v>(HashIdentity.Structural))

let apply concestor (d: Dictionary<_, _>) op key value newValue =
  match op with
  | Add ->
      d.[key] <- value
      let noValue = Unchecked.defaultof<_>
      Delta(Remove, key, value, noValue, concestor)
  | Remove ->
      let _ = d.Remove key
      let noValue = Unchecked.defaultof<_>
      Delta(Add, key, value, noValue, concestor)
  | Replace ->
      d.[key] <- newValue
      Delta(Replace, key, newValue, value, concestor)

let rec ossifyCPS concestor k =
  match concestor.Ancestry with
  | HashTable s -> k s
  | Delta(op, key, value, newValue, relative) ->
      ossifyCPS relative (fun (d: Dictionary<_, _>) ->
        relative.Ancestry <-
          apply concestor d op key value newValue
        // OPTIM: We only need to set this once.
        concestor.Ancestry <- HashTable d
        k d)

let ossify concestor = ossifyCPS concestor id

let locked concestor f =
  lock concestor.LockObj (fun () -> f(ossify concestor))

let add key value concestor =
  locked concestor (fun d ->
    let curr = { concestor with Ancestry = HashTable d }
    let ancestry =
      let mutable oldValue = Unchecked.defaultof<_>
      if d.TryGetValue(key, &oldValue) then
        d.[key] <- value
        Delta(Replace, key, value, oldValue, curr)
      else
        d.[key] <- value
        Delta(Remove, key, value, Unchecked.defaultof<_>, curr)
    concestor.Ancestry <- ancestry
    curr)

let remove key concestor =
  locked concestor (fun d ->
    let mutable oldValue = Unchecked.defaultof<_>
    if d.TryGetValue(key, &oldValue) then
      let _ = d.Remove key
      let curr = { concestor with Ancestry = HashTable d }
      let noValue = Unchecked.defaultof<_>
      concestor.Ancestry <- Delta(Add, key, oldValue, noValue, curr)
      curr
    else concestor)

let containsKey key concestor =
  locked concestor (fun d -> d.ContainsKey key)

let find key concestor =
  locked concestor (fun d -> d.[key])

let tryFind key concestor =
  locked concestor (fun d ->
    let mutable value = Unchecked.defaultof<_>
    if d.TryGetValue(key, &value) then
      Some value
    else
      None)

let count concestor =
  locked concestor (fun d -> d.Count)

let toArray concestor = locked concestor Array.ofSeq

let fold f a concestor =
  toArray concestor
  |> Array.fold (fun a (KeyValue(k, v)) -> f a k v) a

let isEmpty concestor =
  count concestor = 0

let ofSeq kvs =
  let d = Dictionary(HashIdentity.Structural)
  for k, v in kvs do
    d.[k] <- v
  create d