@@ -0,0 +1,123 @@
1+module Data.SortedBag
2+
3+import Control.Isomorphism
4+import Control.Pipeline
5+import Data.Chain
6+import Data.Combinators
7+import Data.PosNat
8+import Data.SortedMap
9+
10+%default total
11+%access export
12+
13+data SortedBag k = BagWrapper (SortedMap k PosNat)
14+
15+||| Remove the `BagWrapper` from a bag (the required proof from `PosNat` makes
16+||| this safe)
17+unwrapBag : SortedBag k -> SortedMap k PosNat
18+unwrapBag (BagWrapper m) = m
19+
20+||| An `Isomorphism` between a wrapped and unwrapped bag
21+wrappedBagIso : Iso (SortedMap k PosNat) (SortedBag k)
22+wrappedBagIso = MkIso BagWrapper unwrapBag (\(BagWrapper _) => Refl) (\_ => Refl)
23+
24+||| A type for converting a function from an operation over SortedMaps to an
25+||| operation over SortedBags
26+OnMapsT : Nat -> Type -> Type
27+OnMapsT n k = (EndoChain n $ SortedMap k PosNat) -> (EndoChain n $ SortedBag k)
28+
29+||| Apply an n-ary function over the wrapped `SortedMap`.
30+onMaps : OnMapsT n k
31+onMaps = withIso wrappedBagIso
32+
33+||| Special case for n=1 for `onMaps`
34+onMap : OnMapsT 1 k
35+onMap = onMaps {n=1}
36+
37+||| An empty bag.
38+empty : Ord k => SortedBag k
39+empty = BagWrapper empty
40+
41+-- Keep in mind:
42+-- Idris> the Nat $ sum Nothing
43+-- 0 : Nat
44+-- Idris> the Nat $ sum $ Just 10
45+-- 10 : Nat
46+
47+||| The type of `insert` and the common core of most functions that update
48+||| the contests of a Bag.
49+UpdateBag : Type -> Type
50+UpdateBag k = k -> SortedBag k -> SortedBag k
51+
52+||| Apply a function `(Nat -> Nat)` to a given argument
53+alterCounts : (Nat -> Nat) -> UpdateBag k
54+alterCounts f k = onMap $ \m => let
55+ n = sum $ fst <$> lookup k m
56+ n' = f n
57+ in m |> if n == n'
58+ then id -- No change before and after.
59+ else case n' of
60+ Z => delete k
61+ n'@(S _) => insert k (n' ** ItIsSucc)
62+
63+||| Count the number of occurances
64+count : k -> SortedBag k -> Nat
65+count k (BagWrapper m) = sum $ fst <$> lookup k m
66+
67+||| Is a count non-zero?
68+contains : k -> SortedBag k -> Bool
69+contains = isZero ... count
70+
71+||| insert a given number of items
72+insertN : Nat -> UpdateBag k
73+insertN n = alterCounts (plus n)
74+
75+||| insert a single item.
76+insert : UpdateBag k
77+insert = alterCounts S
78+
79+||| delete a given number of items
80+deleteN : Nat -> UpdateBag k
81+deleteN n = alterCounts (`minus` n)
82+
83+||| delete 1 item (equivalent to `deleteN 1`)
84+delete1 : UpdateBag k
85+delete1 = deleteN 1
86+
87+||| Delete all ocurances
88+delete : UpdateBag k
89+delete = alterCounts $ const Z
90+
91+||| A single item (more efficient than `insert x empty`)
92+singleton : Ord k => k -> SortedBag k
93+singleton k = BagWrapper $ Data.SortedMap.insert k one empty
94+-- So that we don't need to do a fetch first
95+
96+-- The reason for rewriting this file, because now it is
97+||| O(1). Are there any items? The proof in PosNat ensures that there are no
98+||| zero-counted items.
99+null : SortedBag k -> Bool
100+null (BagWrapper m) = null m
101+
102+||| Convert a list to a Sorted Bag
103+fromList : Ord k => List k -> SortedBag k
104+fromList = foldr insert empty
105+
106+||| Convert a Sorted Bag to a List, keeping it sorted.
107+toList : SortedBag k -> List k
108+toList (BagWrapper m) = concat $ (\(x, (n ** _)) => replicate n x) <$> toList m
109+
110+-- TODO: more efficient, implementation also for Map and Set
111+Foldable SortedBag where
112+ foldl f zero = foldl f zero . Data.SortedBag.toList
113+ foldr f zero = foldr f zero . Data.SortedBag.toList
114+
115+||| Merge two bags
116+union : EndoChain 2 (SortedBag k)
117+union = onMaps {n=2} merge
118+
119+Semigroup (SortedBag a) where
120+ (<+>) = union
121+
122+Ord a => Monoid (SortedBag a) where
123+ neutral = empty