45 lines
1.4 KiB
Haskell
45 lines
1.4 KiB
Haskell
{-# OPTIONS_GHC -Wno-unused-imports #-}
|
|
module IntMapHelp
|
|
( module Data.Strict.IntMap
|
|
, newKey
|
|
, insertNewKey
|
|
, insertWithNewKeys
|
|
, unsafeSwapKeys
|
|
, safeSwapKeys
|
|
, findIndex
|
|
) where
|
|
import LensHelp
|
|
--import Data.List hiding (foldr,insert)
|
|
import Data.Strict.IntMap
|
|
import Data.Strict.Containers.Lens
|
|
{- | Find a key value one higher than any key in the map, or zero if the map is
|
|
- empty -}
|
|
newKey :: IntMap a -> Int
|
|
newKey = maybe 0 ((+ 1) . fst) . lookupMax
|
|
{- | Insert an element with some new key. -}
|
|
insertNewKey :: a -> IntMap a -> IntMap a
|
|
insertNewKey x m = case lookupMax m of
|
|
Nothing -> singleton 0 x
|
|
Just (k,_) -> insert (k+1) x m
|
|
{- | Insert a list of values with new keys -}
|
|
insertWithNewKeys :: [a] -> IntMap a -> IntMap a
|
|
insertWithNewKeys xs m = union m $ fromList $ zip [newKey m..] xs
|
|
{- | Swaps two keys. Unsafe: assumes both keys exist. -}
|
|
unsafeSwapKeys :: Int -> Int -> IntMap a -> IntMap a
|
|
unsafeSwapKeys i j m = insert i (m ! j) $ insert j (m ! i) m
|
|
|
|
safeSwapKeys :: Int -> Int -> IntMap a -> IntMap a
|
|
safeSwapKeys i j m = m
|
|
& at i .~ (m ^? ix j)
|
|
& at j .~ (m ^? ix i)
|
|
|
|
-- ideally this should short circuit, I am not sure whether it does or not
|
|
-- though
|
|
findIndex :: (a -> Bool) -> IntMap a -> Maybe Int
|
|
findIndex t = foldrWithKey f Nothing
|
|
where
|
|
f _ _ (Just k) = Just k
|
|
f k x _
|
|
| t x = Just k
|
|
| otherwise = Nothing
|