----------------------------- MODULE ParallelSort -----------------------------
(***************************************************************************)
(* A TLA+ model of the fork/join structure shared by ParallelMerge        *)
(* (sort/parallelmerge.go) and ParallelQuick (sort/parallelquick.go).      *)
(*                                                                         *)
(* Both sorts recursively split an array range into two halves, sort them  *)
(* in two goroutines over DISJOINT index sub-ranges, join on a WaitGroup,  *)
(* and (for merge) combine. The correctness of the *parallelization* — as  *)
(* argued on paper in docs/verification.md — reduces to two claims:        *)
(*                                                                         *)
(*   1. No data race: no two concurrently-writing tasks touch overlapping  *)
(*      array indices.                                                      *)
(*   2. Termination: the join structure always completes (no deadlock).    *)
(*                                                                         *)
(* This model checks both EXHAUSTIVELY over the whole recursion tree. We   *)
(* model each task's index range as the resource it writes. A parent only  *)
(* writes its full range (the merge step) AFTER both children have joined  *)
(* (reached "done"); until then it is "waiting" and writes nothing. This is *)
(* the WaitGroup fence. Remove that fence (let a parent merge while its     *)
(* children still run) and TLC finds a data race — see the README.         *)
(*                                                                         *)
(* Modeling the parent's post-join merge as writing the whole [lo,hi) range *)
(* is the *stronger* case (merge sort). Quicksort does no work after the    *)
(* join, so if this race-free model holds, quicksort's is race-free too.    *)
(***************************************************************************)
EXTENDS Naturals, FiniteSets

CONSTANTS
    N,          \* array length (indices 0..N-1)
    Threshold   \* ranges of size <= Threshold are sorted sequentially (leaves)

\* The recursion tree: node id 1 owns [0,N); a non-leaf id owns [lo,hi) and its
\* children 2*id, 2*id+1 own the two halves. Ranges of size <= Threshold are
\* leaves. The set of nodes is finite because ranges strictly shrink.
RECURSIVE TreeFrom(_, _, _)
TreeFrom(id, lo, hi) ==
    IF hi - lo <= Threshold
    THEN { [id |-> id, lo |-> lo, hi |-> hi, leaf |-> TRUE] }
    ELSE LET mid == (lo + hi) \div 2
         IN  { [id |-> id, lo |-> lo, hi |-> hi, leaf |-> FALSE] }
             \union TreeFrom(2*id, lo, mid)
             \union TreeFrom(2*id + 1, mid, hi)

Tree    == TreeFrom(1, 0, N)
NodeIds == { t.id : t \in Tree }
Node(i) == CHOOSE t \in Tree : t.id = i     \* the record for id i
IsLeaf(i)  == Node(i).leaf
Left(i)  == 2 * i
Right(i) == 2 * i + 1

\* Two half-open index ranges are disjoint iff one ends at or before the other
\* begins.
Disjoint(i, j) == \/ Node(i).hi <= Node(j).lo
                  \/ Node(j).hi <= Node(i).lo

VARIABLE phase   \* phase[i] \in {"unstarted","ready","waiting","merging","done"}
vars == << phase >>

\* A task is *writing* its range while it computes: a leaf sorting its small
\* range, or an internal node performing its post-join merge.
Writing == { i \in NodeIds : phase[i] = "merging" }

TypeOK ==
    phase \in [NodeIds -> {"unstarted", "ready", "waiting", "merging", "done"}]

Init ==
    \* Only the root starts ready; children become ready when their parent forks.
    phase = [i \in NodeIds |-> IF i = 1 THEN "ready" ELSE "unstarted"]

\* Fork: a ready internal node spawns its two children (disjoint halves) and
\* waits for them (WaitGroup.Add(2); go ...; go ...; wg.Wait()).
Fork(i) ==
    /\ phase[i] = "ready"
    /\ ~IsLeaf(i)
    /\ phase' = [phase EXCEPT ![i]        = "waiting",
                              ![Left(i)]  = "ready",
                              ![Right(i)] = "ready"]

\* Leaf: a ready leaf sorts its range sequentially (enters the writing phase).
Leaf(i) ==
    /\ phase[i] = "ready"
    /\ IsLeaf(i)
    /\ phase' = [phase EXCEPT ![i] = "merging"]

\* Join+merge: a waiting node whose BOTH children are done may now write its
\* full range. This guard is the WaitGroup fence -- the whole point of the check.
Merge(i) ==
    /\ phase[i] = "waiting"
    /\ phase[Left(i)]  = "done"
    /\ phase[Right(i)] = "done"
    /\ phase' = [phase EXCEPT ![i] = "merging"]

\* A writing task finishes.
Finish(i) ==
    /\ phase[i] = "merging"
    /\ phase' = [phase EXCEPT ![i] = "done"]

AllDone == phase[1] = "done"   \* the root finishes only after its whole subtree

Next ==
    \/ \E i \in NodeIds : Fork(i) \/ Leaf(i) \/ Merge(i) \/ Finish(i)
    \/ (AllDone /\ UNCHANGED vars)   \* stutter when finished (no false deadlock)

Spec ==
    /\ Init /\ [][Next]_vars
    /\ \A i \in NodeIds : WF_vars(Fork(i) \/ Leaf(i) \/ Merge(i) \/ Finish(i))

-----------------------------------------------------------------------------
\* Properties (see ParallelSort.cfg).

\* SAFETY / no data race: all tasks writing at the same time own disjoint ranges.
NoDataRace ==
    \A i, j \in Writing : (i /= j) => Disjoint(i, j)

\* LIVENESS / termination: the fork/join always completes (no deadlock).
Terminates == <>AllDone
=============================================================================