waterdeep/src/Waterdeep/Util.hs

92 lines
3.0 KiB
Haskell

{-# LANGUAGE LambdaCase #-}
{-# LANGUAGE RankNTypes #-}
module Waterdeep.Util
( deleteAt
, on
, nubOn
, sortOn
, countOf
, singular
, mif
, joinStrings
, (<$$>)
, whenM
, unlessM
, nextIntKey
, intAtNext
, walkIntMap
) where
import Control.Applicative
import Control.Monad.State
import Data.List
import Data.Function
import Data.Maybe
import Data.Monoid
import Lens.Family2
import Lens.Family2.State
import Lens.Family2.Stock
import qualified Data.IntMap as IM
deleteAt :: Int -> [a] -> [a]
deleteAt n l = left ++ drop 1 right
where (left, right) = splitAt n l
nubOn :: Eq b => (a -> b) -> [a] -> [a]
nubOn f = nubBy ((==) `on` f)
sortOn :: Ord b => (a -> b) -> [a] -> [a]
sortOn f = sortBy (compare `on` f)
countOf :: Num r => FoldLike (Sum r) a a' b b' -> (b -> Bool) -> a -> r
countOf l f = getSum . views l (\b -> if f b then Sum 1 else Sum 0)
-- |Turn a traversal into a lens. This is well-defined if and only if the traversal
-- always refers to exactly one element. With multiple elements, setting applies to
-- all and getting returns just the first value. With zero elements, setting is a
-- no-op and getting results in a runtime error.
singular :: Traversal a a' b b' -> Lens a a' b b'
singular t f b = (\a' -> b & t .~ a') `fmap` (f (unJust (firstOf t b)))
where unJust (Just x) = x
unJust Nothing = error "singular: empty traversal"
mif :: Monoid m => Bool -> m -> m
mif c m = if c then m else mempty
joinStrings :: [String] -> String
joinStrings [] = "nothing"
joinStrings [x] = x
joinStrings [x,y] = x ++ " and " ++ y
joinStrings [x,y,z] = x ++ ", " ++ y ++ ", and " ++ z
joinStrings (x:xs) = x ++ ", " ++ joinStrings xs
-- Like (<$>), but with the same fixity and precedence as ($)
(<$$>) :: Applicative f => (a -> b) -> f a -> f b
(<$$>) = (<$>)
infixr 0 <$$>
-- Like when/unless, except that the condition is an action yielding a Bool
-- rather than a pure Bool. Avoids the need for spurious condition variables.
whenM, unlessM :: Monad m => m Bool -> m () -> m ()
whenM mc ma = do { c <- mc; when c ma }
unlessM mc ma = do { c <- mc; unless c ma }
-- Returns the next available key in an IntMap, starting from 1 if the map is empty.
nextIntKey :: IM.IntMap a -> IM.Key
nextIntKey im = if IM.null im then 1 else 1 + fst (IM.findMax im)
-- This isn't really a lens since it doesn't refer to a constant location.
-- You don't get back what you put in, and setting twice has a different
-- effect than setting once. Nonetheless, I find it useful. It works like
-- intAt except that the (implied) key parameter always refers to the next
-- available key.
intAtNext :: Applicative f => (Maybe b -> f (Maybe b)) -> IM.IntMap b -> f (IM.IntMap b)
intAtNext f im = maybe im (flip (IM.insert (nextIntKey im)) im) <$> f Nothing
walkIntMap :: MonadState a m => Lens' a (IM.IntMap b)
-> (IM.Key -> Lens' a (Maybe b) -> m r) -> m ()
walkIntMap l f = use l >>= \im -> unless (IM.null im) $ go (fst $ IM.findMin im) where
go k = f k (l . intAt k) >> use l >>= maybe (return ()) (go . fst) . IM.lookupGT k