Simpler method of making multiple changes to a function in a single step

Hi,

I have a set of nodes N, a set of resources R and a function "register" that maps N to R. The algorithm is a resource allocation algorithm that must assign the resources of R evenly across the nodes N.

So if N is { "n1", "n2", "n3"} and R is {"r1", "r2", "r3", "r4" } then once allocation has taken place a valid value for register would be:

[n1 |-> <<"r4", "r1">>, n2 |-> <<"r2">>, n3 |-> <<"r3">>]

I want to set the values of register in a single step and I have managed it, though the formula seems overly complex and I wonder if there is a simpler way of doing that would help me also gain more insight into TLA+.

I have isolated the allocation logic into a toy spec as follows:

EXTENDS Integers, FiniteSets, Sequences, TLC

CONSTANT R, N

VARIABLE register

Init ==

HasMinResources(counts, nd) ==

\A x \in N : counts[nd] <= counts[x]

Allocate ==

LET newRegister == [n \in N |-> << >>]

counts == [n \in N |-> 0]

al[pendingRes \in SUBSET R, assignCount \in [N -> Nat]] ==

LET n == CHOOSE nd \in N : HasMinResources(assignCount, nd)

r == LET int == R \intersect pendingRes

IN CHOOSE x \in int : TRUE

IF Cardinality(pendingRes) = 0 THEN newRegister

ELSE

LET remaining == pendingRes \ { r }

newAssignCount == [assignCount EXCEPT ![n] = @ + 1]

IN [al[remaining, newAssignCount] EXCEPT ![n] = Append(@, r)]

IN al[R, counts]

Rebalance ==

/\ register' = Allocate

/\ PrintT(register')

(* ignore the spec, I just wanted to run the Rebalance action once *)

Spec == Init /\ [][Rebalance]_register

Notes:

- I made Allocate recursive as that is the only way I could figure out making all the changes to register in a single step.

- I did the intersect so that I could use CHOOSE. Else it complained that is was an unbounded CHOOSE so I figured if I did an intersect with R then it would be interpretted as bounded.

Any insights or suggestions would be great.

Thanks

Jack