0 | ||| Specialised set of `Fin`s.
  1 | ||| Isomorphic to a vector of appropriate size of booleans.
  2 | |||
  3 | ||| This should be a replacement for `SortedSet (Fin n)` working better at least in elaborator scripts.
  4 | ||| This can be implemented as simple `Integer` setting and clearing bits, but this seems to work worse in elaborator scripts.
  5 | module Data.Fin.Set
  6 |
  7 | import Data.Bits
  8 | import Data.Fin.Map
  9 | import Data.SortedMap
 10 | import Data.SortedSet
 11 | import Data.String
 12 | import Data.Vect
 13 |
 14 | %default total
 15 |
 16 | export
 17 | data FinSet : Nat -> Type where
 18 |   MkFS : Integer -> FinSet n
 19 |
 20 | mask : (n : Nat) -> Integer
 21 | mask n = (1 `shiftL` n) - 1
 22 |
 23 | %inline
 24 | unFS : FinSet n -> Integer
 25 | unFS (MkFS bs) = bs
 26 |
 27 | %inline
 28 | wrap : (Integer -> Integer) -> FinSet n -> FinSet n
 29 | wrap f $ MkFS x = MkFS $ f x
 30 |
 31 | %inline
 32 | wrap2 : (Integer -> Integer -> Integer) -> FinSet n -> FinSet n -> FinSet n
 33 | wrap2 f (MkFS l) (MkFS r) = MkFS $ f l r
 34 |
 35 | export %inline
 36 | empty : FinSet n
 37 | empty = MkFS zeroBits
 38 |
 39 | export %inline
 40 | full : FinSet n
 41 | full = MkFS oneBits
 42 |
 43 | export %inline
 44 | complement : FinSet n -> FinSet n
 45 | complement = wrap $ Bits.complement
 46 |
 47 | export %inline
 48 | insert : Fin n -> FinSet n -> FinSet n
 49 | insert i = wrap (`setBit` finToNat i)
 50 |
 51 | export %inline
 52 | insert' : FinSet n -> Fin n -> FinSet n
 53 | insert' = flip insert
 54 |
 55 | export %inline
 56 | delete : Fin n -> FinSet n -> FinSet n
 57 | delete i = wrap (`clearBit` finToNat i)
 58 |
 59 | export %inline
 60 | delete' : FinSet n -> Fin n -> FinSet n
 61 | delete' = flip delete
 62 |
 63 | export %inline
 64 | contains : Fin n -> FinSet n -> Bool
 65 | contains i = (`testBit` finToNat i) . unFS
 66 |
 67 | export %inline
 68 | contains' : FinSet n -> Fin n -> Bool
 69 | contains' = flip contains
 70 |
 71 | export %inline
 72 | fromFoldable : Foldable f => f (Fin n) -> FinSet n
 73 | fromFoldable = foldl insert' $ empty {n}
 74 |
 75 | export %inline
 76 | fromList : List (Fin n) -> FinSet n
 77 | fromList = fromFoldable
 78 |
 79 | export %inline
 80 | toList : {n : _} -> FinSet n -> List (Fin n)
 81 | toList {n=Z}   = const []
 82 | toList {n=S n} = go n last [] . unFS where
 83 |   go : Nat -> Fin (S n) -> List (Fin $ S n) -> Integer -> List (Fin $ S n)
 84 |   go _ FZ        acc i = if testBit i 0 then FZ :: acc else acc
 85 |   go n k@(FS k') acc i = let next = if testBit i n then k :: acc else acc in go (n `minus` 1) (assert_smaller k $ believe_me k') next i
 86 |
 87 | export
 88 | traverse_ : Monad m => {n : _} -> (Fin n -> m b) -> FinSet n -> m ()
 89 | traverse_ {n=Z}       _ = const $ pure ()
 90 | traverse_ {n=n@(S _)} f = go FZ . (mask n .&.) . unFS where
 91 |   go : Fin n -> Integer -> m ()
 92 |   go _   0 = pure ()
 93 |   go idx i = do
 94 |     let %inline next : m ()
 95 |         next = go (believe_me $ FS idx) (assert_smaller i $ i `shiftR` 1)
 96 |     if testBit i 0 then ignore (f idx) >> next else next
 97 |
 98 | public export %inline
 99 | for_ : Monad m => {n : _} -> FinSet n -> (Fin n -> m b) -> m ()
100 | for_ = flip traverse_
101 |
102 | export
103 | foldl : {n : _} -> (a -> Fin n -> a) -> a -> FinSet n -> a
104 | foldl {n=Z}       _ ini = const ini
105 | foldl {n=n@(S _)} f ini = go FZ ini . unFS where
106 |   go : Fin n -> a -> Integer -> a
107 |   go _   curr 0 = curr
108 |   go idx curr i = let next = if testBit i 0 then f curr idx else curr in go (believe_me $ FS idx) next (assert_smaller i $ i `shiftR` 1)
109 |
110 | public export
111 | foldMap : Monoid m => {n : _} -> (Fin n -> m) -> FinSet n -> m
112 | foldMap f = foldl (\c => (c <+>) . f) neutral
113 |
114 | public export %inline
115 | (.asList) : {n : _} -> FinSet n -> List (Fin n)
116 | (.asList) = toList
117 |
118 | public export %inline
119 | size : {n : _} -> FinSet n -> Nat
120 | size = length . toList {n} -- this implementation is to make `asVect` work seamlessly
121 |
122 | public export %inline
123 | (.size) : {n : _} -> FinSet n -> Nat
124 | (.size) = size {n}
125 |
126 | public export %inline
127 | (.asVect) : {n : _} -> (fs : FinSet n) -> Vect fs.size (Fin n)
128 | (.asVect) fs = fromList $ toList fs
129 |
130 | export %inline
131 | union : FinSet n -> FinSet n -> FinSet n
132 | union = wrap2 (.|.)
133 |
134 | export %inline
135 | difference : FinSet n -> FinSet n -> FinSet n
136 | difference (MkFS l) (MkFS r) = MkFS $ l .&. complement r
137 |
138 | export %inline
139 | symDifference : FinSet n -> FinSet n -> FinSet n
140 | symDifference = wrap2 xor
141 |
142 | export %inline
143 | intersection : FinSet n -> FinSet n -> FinSet n
144 | intersection = wrap2 (.&.)
145 |
146 | export %inline
147 | leftMost : {n : _} -> FinSet n -> Maybe $ Fin n
148 | leftMost = go n . unFS where
149 |   go : (n : Nat) -> Integer -> Maybe $ Fin n
150 |   go _     0 = Nothing
151 |   go 0     _ = Nothing
152 |   go (S n) x = if testBit x 0 then Just FZ else FS <$> go n (x `shiftR` 1)
153 |
154 | export %inline
155 | rightMost : {n : _} -> FinSet n -> Maybe $ Fin n
156 | rightMost = go n . unFS where
157 |   go : (n : Nat) -> Integer -> Maybe $ Fin n
158 |   go _     0 = Nothing
159 |   go 0     _ = Nothing
160 |   go (S n) x = if testBit x n then Just last else weaken <$> go n x
161 |
162 | export
163 | {n : _} -> Show (FinSet n) where
164 |   showPrec d s = showCon d "fromList" $ showArg $ toList s
165 |
166 | export
167 | {n : _} -> Eq (FinSet n) where
168 |   (==) = (==) `on` (mask n .&.) . unFS
169 |
170 | export
171 | {n : _} -> Ord (FinSet n) where
172 |   compare = comparing $ (mask n .&.) . unFS
173 |
174 | export
175 | Semigroup (FinSet n) where
176 |   (<+>) = union {n}
177 |
178 | export
179 | Monoid (FinSet n) where
180 |   neutral = empty {n}
181 |
182 | export
183 | {n : _} -> Interpolation (Fin n) => Interpolation (FinSet n) where
184 |   interpolate = ("{" ++) . (++ "}") . joinBy ", " . map interpolate . Fin.Set.toList {n}
185 |
186 | export
187 | null : {n : _} -> FinSet n -> Bool
188 | null (MkFS x) = mask n .&. x == 0
189 |
190 | namespace SortedMap
191 |
192 |   export
193 |   keySet : {n : _} -> SortedMap (Fin n) v -> FinSet n
194 |   keySet = fromFoldable . map fst . SortedMap.toList
195 |
196 |   export
197 |   keySetSize : (m : SortedMap (Fin n) v) -> (keySet m).size = length (SortedMap.toList m)
198 |   keySetSize _ = believe_me $ Refl {x=Z}
199 |
200 | namespace FinMap
201 |
202 |   export
203 |   keySet : {n : _} -> FinMap n v -> FinSet n
204 |   keySet = fromFoldable . map fst . kvList
205 |
206 |   export
207 |   keySetSize : (m : FinMap n v) -> (keySet m).size = length (kvList m)
208 |   keySetSize _ = believe_me $ Refl {x=Z}
209 |
210 | namespace SortedSet
211 |
212 |   export %inline
213 |   toSortedSet : {n : _} -> FinSet n -> SortedSet (Fin n)
214 |   toSortedSet = fromList . toList
215 |
216 | export
217 | singleton : Fin n -> FinSet n
218 | singleton = MkFS . setBit zeroBits . finToNat
219 |
220 | export
221 | fit : {n : _} -> FinSet n -> FinSet m
222 | fit (MkFS x) = MkFS $ mask n .&. x
223 |
224 | export %inline
225 | weakenToSuper : {i : _} -> FinSet (finToNat i) -> FinSet n
226 | weakenToSuper = fit
227 |