Shuffle vector model in TLA+

ShuffleVector.tla

------------------------- MODULE ShuffleVector ------------------------------

LOCAL INSTANCE Naturals
LOCAL INSTANCE Sequences
LOCAL INSTANCE TLC

(* A Shuffle Vector consists of a sequence of up to N elements.
   Adding an element to the shuffle vector pushes an element onto the beginning 
   of the sequence and swaps it with a random element.
 *)

CONSTANT MaxSize

ASSUME MaxSize \in Nat
ASSUME MaxSize > 0

(* Finite sequence definition from https://learntla.com/tla/functions/ *)
LOCAL FiniteSeq(S, n) == UNION {[1..m -> S] : m \in 1..n}

(* A sequence has no repeated elements iff it is injective *)
LOCAL NoRepeatedElements(SV) == 
  ~\E x \in DOMAIN(SV), y \in DOMAIN(SV):
    x /= y /\ SV[x] = SV[y]
    
(* A shuffle vector is a sequence of up to MaxSize elements, and may be empty.
   It should contain no repeated elements. *) 
ShuffleVectors(Elements) == 
 { s \in FiniteSeq(Elements,MaxSize) \union { <<>> } : NoRepeatedElements( s ) }
   
(* Is the shuffle vector empty? *)
Empty(SV) == Len( SV ) = 0

(* Is the shuffle vector full? *)
Full(SV) == Len( SV ) = MaxSize

(* What is the next element to pop from the shuffle vector? *)
NextElement(SV) == Head( SV )

(* Return the shuffle vector obtained by popping the next element *)
Pop(SV) == 
    IF Len( SV ) > 0 THEN 
        Tail( SV )
    ELSE
        Assert( FALSE, "Pop from empty shuffle vector." )

(* Range definition from https://learntla.com/tla/sets/, necessary
   to recover the type of Elements for the SV for use in the recursive
   function below. *)
LOCAL Range(T) == { T[x] : x \in DOMAIN T }
        
(* Swap element 1 with a randomly-chosen element, including itself.
   TLC won't explore *all* possibilities for the random element, just one
   per state, so the state graph won't be complete.
   I don't know how to fix this other than take in an additional parameter
   that can go into an existential quantifier.
   One option would be to return all the possible successor states and
   let the action choose one, I guess (but not CHOOSE.) *)
LOCAL RandomSwap(SV) ==
    LET idx == RandomElement(DOMAIN SV)
        Elements == Range(SV)
        Swap[position \in Nat, left,right \in Seq(Elements)] == 
            IF position = 1 THEN
                <<Head(right)>> \o Tail(left) \o <<Head(left)>> \o Tail(right)
            ELSE
                Swap[ position - 1, Append(left, Head(right)), Tail(right)]
    IN IF idx = 1 THEN
        SV \* Special-casing this seemed easier than working out how to handle zero-length "left" in the recursive function
    ELSE
        Swap[idx - 1, <<Head(SV)>>, Tail(SV)]
   
(* Return the shuffle vector obtained by pushing X. *)
Push(SV,X) == 
    IF Len( SV ) < MaxSize THEN
        RandomSwap( <<X>> \o SV ) 
    ELSE
        Assert( FALSE, "Push into full shuffle vector." )     
   
=============================================================================

TestShuffleVector.tla

------------------------- MODULE TestShuffleVector -------------------------

EXTENDS Naturals, FiniteSets, Sequences

CONSTANT SpanSize

INSTANCE ShuffleVector WITH MaxSize <- SpanSize

VARIABLE sv, inuse

SpanIndexes == 0..(SpanSize-1)

SV == ShuffleVectors(SpanIndexes)

Init == 
    /\ sv \in SV
    /\ Full(sv)
    /\ inuse = {}

TypeInvariant ==
    /\ sv \in SV
    /\ inuse \in SUBSET SpanIndexes
    
Alloc == 
    /\ ~Empty( sv )
    /\ inuse' = inuse \union { NextElement( sv ) }
    /\ sv' = Pop( sv )

Dealloc ==
    \E x \in SpanIndexes :
        /\ x \in inuse 
        /\ inuse' = inuse \ { x }
        /\ sv' = Push( sv, x )
        
Next == Alloc \/ Dealloc

Spec == Init /\ [Next]_<<sv,inuse>>

NoMemoryLeak == Len( sv ) + Cardinality( inuse ) = SpanSize

=============================================================================