fold operations for bags

Author: muenchnerkindl

modules/BagsExt.tla

@@ -4,12 +4,13 @@
 (***************************************************************************)
 
 LOCAL INSTANCE Bags
-LOCAL INSTANCE Naturals
+LOCAL INSTANCE Integers
 LOCAL INSTANCE Folds
 
 BagAdd(B, x) ==
    (************************************************************************)
    (* Add an element x to bag B.                                           *)
+   (* Same as B (+) SetToBag({x}).                                         *)
    (************************************************************************)
    IF x \in DOMAIN B
    THEN [e \in DOMAIN B |-> IF e=x THEN B[e]+1 ELSE B[e]]
@@ -18,11 +19,43 @@ BagAdd(B, x) ==
 BagRemove(B,x) ==
    (************************************************************************)
    (* Remove an element x from bag B.                                      *)
+   (* Same as B (-) SetToBag({x}).                                         *)
    (************************************************************************)
    IF x \in DOMAIN B
    THEN IF B[x] = 1
         THEN [e \in DOMAIN B \ {x} |-> B[e]]
         ELSE [e \in DOMAIN B |-> IF e=x THEN B[e]-1 ELSE B[e]]
    ELSE B
 
+FoldBag(op(_,_), base, B) ==
+   (************************************************************************)
+   (* Fold op over all elements of bag B, starting with value base.        *)
+   (*                                                                      *)
+   (* Example:                                                             *)
+   (*    FoldBag(LAMBDA x,y : x+y,                                         *)
+   (*            0,                                                        *)
+   (*            (1 :> 2) @@ (2 :> 1) @@ (3 :> 2)) = 10                    *)
+   (*    The third argument represents the bag {| 1, 1, 2, 3, 3 |}.        *)
+   (************************************************************************)
+   LET repl(x) ==
+         (* replicate operation op for all occurrences of x in B *)
+       LET it[n \in Nat \ {0}] ==
+           IF n = 1 THEN x ELSE op(x, it[n-1])
+       IN  it[B[x]]
+   IN  MapThenFoldSet(op, base, repl, 
+                      LAMBDA S : CHOOSE x \in S : TRUE,
+                      DOMAIN B)
+
+SumBag(B) ==
+   (************************************************************************)
+   (* Compute the sum of the elements of B.                                *)
+   (************************************************************************)
+   FoldBag(LAMBDA x,y : x+y, 0, B)
+
+ProductBag(B) ==
+   (************************************************************************)
+   (* Compute the product of the elements of B.                            *)
+   (************************************************************************)
+   FoldBag(LAMBDA x,y : x*y, 1, B)
+
 =============================================================================

tests/BagsExtTests.tla

@@ -0,0 +1,31 @@
+---- CONFIG BagsExtTests ----
+INIT Init
+NEXT Next
+=====
+
+---------------------------- MODULE BagsExtTests ----------------------------
+EXTENDS Integers, TLC, Bags, BagsExt
+
+ASSUME LET T == INSTANCE TLC IN T!PrintT("BagsExtTests")
+
+-----------------------------------------------------------------------------
+
+ASSUME LET B == SetToBag({1,2}) IN BagAdd(B,2) = B (+) SetToBag({2})
+ASSUME LET B == SetToBag({1,2}) IN BagAdd(B,3) = B (+) SetToBag({3})
+ASSUME LET B == SetToBag({1,2}) IN BagRemove(B,2) = B (-) SetToBag({2})
+ASSUME LET B == SetToBag({1,2}) IN 
+           /\ BagRemove(B,3) = B (-) SetToBag({3})
+           /\ BagRemove(B,3) = B
+
+-----------------------------------------------------------------------------
+
+ASSUME FoldBag(LAMBDA x,y : x+y, 0, (1:>2) @@ (2:>1) @@ (3:>2)) = 10
+ASSUME FoldBag(LAMBDA x,y : x+y, 0, (1:>2) @@ (-2:>1)) = 0
+
+ASSUME SumBag((1:>2) @@ (2:>1) @@ (3:>2)) = 10
+ASSUME SumBag((1:>2) @@ (-2:>1)) = 0
+
+ASSUME ProductBag((1:>2) @@ (-2:>1)) = -2
+ASSUME ProductBag((1:>2) @@ (0:>1) @@ (10000:>1)) = 0
+
+=============================================================================