---- MODULE PartitionsStorageSets ----
EXTENDS Naturals, FiniteSets

(*
This model validates only the partition/set logic.
AVL tree lookup and QSBR are abstracted away as the variable `tree`,
which contains ids of currently reachable live partitions.
*)

CONSTANTS
  PartitionIds,
  Threads,
  NoThread,
  KeyCount,
  MinPartitionSize,
  MaxPartitionSize,
  MaxLoad

ASSUME NoThread \notin Threads
ASSUME KeyCount > 0
ASSUME MinPartitionSize > 0
ASSUME MaxPartitionSize >= MinPartitionSize
ASSUME MaxLoad > MaxPartitionSize
ASSUME Cardinality(PartitionIds) >= 3

VARIABLES
  partitions,
  tree,
  locks

vars == <<partitions, tree, locks>>

Keys == 0..(KeyCount - 1)

PartitionRecord ==
  [ lo   : 0..KeyCount,
    hi   : 0..KeyCount,
    dead : BOOLEAN,
    load : 0..MaxLoad ]

KeyInPartition(k, p) ==
  /\ partitions[p].lo <= k
  /\ k < partitions[p].hi

Status(p) ==
  IF partitions[p].dead THEN "DEAD"
  ELSE IF partitions[p].load > MaxPartitionSize THEN "NEED_SPLIT"
  ELSE IF partitions[p].load < MinPartitionSize /\ Cardinality(tree) > 1 THEN "NEED_SHRINK"
  ELSE "ALIVE"

Init ==
  \E p \in PartitionIds:
    /\ tree = {p}
    /\ partitions =
      [ q \in PartitionIds |->
          IF q = p THEN
            [ lo |-> 0,
              hi |-> KeyCount,
              dead |-> FALSE,
              load |-> MinPartitionSize ]
          ELSE
            [ lo |-> 0,
              hi |-> 0,
              dead |-> TRUE,
              load |-> 0 ] ]
    /\ locks = [q \in PartitionIds |-> NoThread]

TypeInvariant ==
  /\ partitions \in [PartitionIds -> PartitionRecord]
  /\ tree \subseteq PartitionIds
  /\ locks \in [PartitionIds -> (Threads \cup {NoThread})]
  /\ \A p \in tree:
      /\ partitions[p].dead = FALSE
      /\ partitions[p].lo < partitions[p].hi
      /\ partitions[p].hi <= KeyCount
  /\ \A p \in PartitionIds \ tree:
      partitions[p].dead = TRUE

UniqueOwnerForEveryKey ==
  \A k \in Keys:
    \E p \in tree:
      /\ KeyInPartition(k, p)
      /\ \A q \in tree:
          KeyInPartition(k, q) => q = p

DeadPartitionsAreOutsideTree ==
  \A p \in PartitionIds:
    partitions[p].dead => p \notin tree

Lock(t, p) ==
  /\ p \in tree
  /\ locks[p] = NoThread
  /\ locks' = [locks EXCEPT ![p] = t]
  /\ UNCHANGED <<partitions, tree>>

Unlock(t, p) ==
  /\ locks[p] = t
  /\ locks' = [locks EXCEPT ![p] = NoThread]
  /\ UNCHANGED <<partitions, tree>>

\* Abstract user mutation of the set stored inside a locked partition.
GrowPartition(t, p) ==
  /\ p \in tree
  /\ locks[p] = t
  /\ partitions[p].load < MaxLoad
  /\ partitions' = [partitions EXCEPT ![p].load = @ + 1]
  /\ UNCHANGED <<tree, locks>>

ShrinkPartitionLoad(t, p) ==
  /\ p \in tree
  /\ locks[p] = t
  /\ partitions[p].load > 0
  /\ partitions' = [partitions EXCEPT ![p].load = @ - 1]
  /\ UNCHANGED <<tree, locks>>

SplitPartition(t, p) ==
  /\ p \in tree
  /\ locks[p] = t
  /\ Status(p) = "NEED_SPLIT"
  /\ partitions[p].lo + 1 < partitions[p].hi
  /\ \E left \in PartitionIds \ tree:
      \E right \in (PartitionIds \ tree) \ {left}:
        \E splitKey \in (partitions[p].lo + 1)..(partitions[p].hi - 1):
          \E leftLoad \in 0..partitions[p].load:
            LET rightLoad == partitions[p].load - leftLoad IN
            /\ rightLoad \in 0..MaxLoad
            /\ leftLoad \in 0..MaxLoad
            /\ partitions' =
              [ partitions EXCEPT
                  ![p].dead = TRUE,
                  ![left] =
                    [ lo |-> partitions[p].lo,
                      hi |-> splitKey,
                      dead |-> FALSE,
                      load |-> leftLoad ],
                  ![right] =
                    [ lo |-> splitKey,
                      hi |-> partitions[p].hi,
                      dead |-> FALSE,
                      load |-> rightLoad ] ]
            /\ tree' = (tree \ {p}) \cup {left, right}
            /\ locks' =
              [ locks EXCEPT
                  ![p] = NoThread,
                  ![left] = NoThread,
                  ![right] = NoThread ]

ShrinkToPrev(t, p) ==
  /\ p \in tree
  /\ locks[p] = t
  /\ Status(p) = "NEED_SHRINK"
  /\ \E prev \in tree \ {p}:
      /\ locks[prev] = t
      /\ partitions[prev].hi = partitions[p].lo
      /\ partitions[prev].load + partitions[p].load <= MaxLoad
      /\ partitions' =
        [ partitions EXCEPT
            ![prev].hi = partitions[p].hi,
            ![prev].load = partitions[prev].load + partitions[p].load,
            ![p].dead = TRUE ]
      /\ tree' = tree \ {p}
      /\ locks' = [locks EXCEPT ![p] = NoThread, ![prev] = NoThread]

ShrinkToNext(t, p) ==
  /\ p \in tree
  /\ locks[p] = t
  /\ Status(p) = "NEED_SHRINK"
  /\ \E next \in tree \ {p}:
      /\ locks[next] = t
      /\ partitions[p].hi = partitions[next].lo
      /\ partitions[next].load + partitions[p].load <= MaxLoad
      /\ partitions' =
        [ partitions EXCEPT
            ![next].lo = partitions[p].lo,
            ![next].load = partitions[next].load + partitions[p].load,
            ![p].dead = TRUE ]
      /\ tree' = tree \ {p}
      /\ locks' = [locks EXCEPT ![p] = NoThread, ![next] = NoThread]

Next ==
  \/ \E t \in Threads, p \in PartitionIds: Lock(t, p)
  \/ \E t \in Threads, p \in PartitionIds: Unlock(t, p)
  \/ \E t \in Threads, p \in PartitionIds: GrowPartition(t, p)
  \/ \E t \in Threads, p \in PartitionIds: ShrinkPartitionLoad(t, p)
  \/ \E t \in Threads, p \in PartitionIds: SplitPartition(t, p)
  \/ \E t \in Threads, p \in PartitionIds: ShrinkToPrev(t, p)
  \/ \E t \in Threads, p \in PartitionIds: ShrinkToNext(t, p)

Spec == Init /\ [][Next]_vars

====