105 lines
3.0 KiB
Haskell
105 lines
3.0 KiB
Haskell
{-# OPTIONS_GHC -Wno-unused-imports #-}
|
|
{-# OPTIONS_GHC -fno-warn-orphans #-}
|
|
|
|
{- |
|
|
WARNING: orphan instances concerning Aeson classes and Data.Strict.IntMap datatypes have been introduced.
|
|
The warnings have been disabled.
|
|
-}
|
|
module IntMapHelp
|
|
-- ( module Data.Strict.IntMap
|
|
(
|
|
module Data.IntMap.Strict,
|
|
newKey,
|
|
insertNewKey,
|
|
insertWithNewKeys,
|
|
unsafeSwapKeys,
|
|
safeSwapKeys,
|
|
findIndex,
|
|
cycleGT,
|
|
cycleLT,
|
|
invertIntMap,
|
|
invertIntMapUnique,
|
|
) where
|
|
|
|
import Control.Applicative
|
|
import Data.Aeson
|
|
import Data.IntMap.Strict
|
|
import qualified Data.Map.Strict as M
|
|
import GHC.Generics
|
|
import LensHelp
|
|
|
|
--import Data.Strict.IntMap
|
|
--import Data.Strict.Containers.Lens
|
|
|
|
--import qualified Data.Strict.IntMap as IM
|
|
--import qualified Data.Strict.IntMap.Autogen.Strict as IM
|
|
--deriving instance Generic (IntMap a)
|
|
|
|
--instance ToJSON1 IntMap where
|
|
-- liftToJSON t tol = liftToJSON to' tol' . toList
|
|
-- where
|
|
-- to' = liftToJSON2 toJSON toJSONList t tol
|
|
-- tol' = liftToJSONList2 toJSON toJSONList t tol
|
|
--
|
|
-- liftToEncoding t tol = liftToEncoding to' tol' . toList
|
|
-- where
|
|
-- to' = liftToEncoding2 toEncoding toEncodingList t tol
|
|
-- tol' = liftToEncodingList2 toEncoding toEncodingList t tol
|
|
--
|
|
--instance ToJSON a => ToJSON (IntMap a) where
|
|
-- toJSON = toJSON1
|
|
-- toEncoding = toEncoding1
|
|
--instance FromJSON a => FromJSON (IntMap a) where
|
|
-- parseJSON = fmap fromList . parseJSON
|
|
|
|
--instance ToJSON a => ToJSON (IntMap a) where
|
|
-- toEncoding = genericToEncoding defaultOptions
|
|
|
|
{- | 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
|
|
|
|
cycleGT :: Int -> IntMap a -> Maybe (Int, a)
|
|
cycleGT i m = lookupGT i m <|> lookupMin m
|
|
|
|
cycleLT :: Int -> IntMap a -> Maybe (Int, a)
|
|
cycleLT i m = lookupLT i m <|> lookupMax m
|
|
|
|
invertIntMap :: Ord a => IntMap a -> M.Map a [Int]
|
|
invertIntMap = foldrWithKey (\k x -> M.insertWith (++) x [k]) mempty
|
|
|
|
invertIntMapUnique :: Ord a => IntMap a -> M.Map a Int
|
|
invertIntMapUnique = foldlWithKey' (\m k x -> M.insert x k m) mempty
|