{-# LANGUAGE MultiParamTypeClasses, FunctionalDependencies,
FlexibleContexts #-}
module Data.Choose.MChoose (
MChoose,
newChoose,
newChoose_,
newListChoose,
newCopyChoose,
copyChoose,
setFirst,
getElem,
setElem,
getPossible,
getSize,
getElems,
setElems,
isValid,
getComplement,
getComplElems,
setNext,
setPrev,
freeze,
unsafeFreeze,
thaw,
unsafeThaw,
unsafeNewListChoose,
unsafeGetElem,
unsafeSetElem,
) where
import Control.Monad
import Control.Monad.ST
import Data.Choose.Base
import Data.Choose.IOBase
class (Monad m) => MChoose c m | c -> m, m -> c where
getPossible :: c -> m Int
getSize :: c -> m Int
newChoose :: Int -> Int -> m c
newChoose_ :: Int -> Int -> m c
unsafeGetElem :: c -> Int -> m Int
unsafeSetElem :: c -> Int -> Int -> m ()
getElems :: c -> m [Int]
setElems :: c -> [Int] -> m ()
unsafeFreeze :: c -> m Choose
unsafeThaw :: Choose -> m c
newListChoose :: (MChoose c m) => Int -> Int -> [Int] -> m c
newListChoose n k is = do
c <- unsafeNewListChoose n k is
valid <- isValid c
when (not valid) $ fail "invalid combination"
return c
{-# INLINE newListChoose #-}
unsafeNewListChoose :: (MChoose c m) => Int -> Int -> [Int] -> m c
unsafeNewListChoose n k is = do
c <- newChoose_ n k
setElems c is
return c
{-# INLINE unsafeNewListChoose #-}
newCopyChoose :: (MChoose c m) => c -> m c
newCopyChoose c = do
n <- getPossible c
k <- getSize c
c' <- newChoose_ n k
copyChoose c' c
return c'
{-# INLINE newCopyChoose #-}
copyChoose :: (MChoose c m) => c -> c -> m ()
copyChoose dst src =
getElems src >>= setElems dst
{-# INLINE copyChoose #-}
setFirst :: (MChoose c m) => c -> m ()
setFirst c = do
k <- getSize c
setElems c [0 .. k-1]
{-# INLINE setFirst #-}
getElem :: (MChoose c m) => c -> Int -> m Int
getElem c i = do
k <- getSize c
when (i < 0 || i >= k) $ fail "getElem: invalid index"
unsafeGetElem c i
{-# INLINE getElem #-}
setElem :: (MChoose c m) => c -> Int -> Int -> m ()
setElem c i x = do
k <- getSize c
when (i < 0 || i >= k) $ fail "getElem: invalid index"
unsafeSetElem c i x
{-# INLINE setElem #-}
isValid :: (MChoose c m) => c -> m Bool
isValid c = do
n <- getPossible c
is <- getElems c
return $! go n (-1) is
where
go _ _ [] = True
go n j (i:is) = i > j && i < n && go n i is
{-# INLINE isValid #-}
getComplement :: (MChoose c m) => c -> m c
getComplement c = do
n <- getPossible c
k <- getSize c
d <- newChoose_ n (n-k)
setElems d =<< getComplElems c
return $! d
{-# INLINE getComplement #-}
getComplElems :: (MChoose c m) => c -> m [Int]
getComplElems c = do
n <- getPossible c
is <- getElems c
return $ go n is 0
where
go n [] j = [j .. n-1]
go n (i:is) j | j == i = go n is (j+1)
| otherwise = [j .. i-1] ++ go n is (i+1)
{-# INLINE getComplElems #-}
setNext :: (MChoose c m) => c -> m Bool
setNext c = do
n <- getPossible c
k <- getSize c
if k > 0
then do
findIncrement (k-1) (n-1) >>=
maybe (return False) (\(i,i') -> do
unsafeSetElem c i (i'+1)
setAscending k (i+1) (i'+2)
return True
)
else
return False
where
findIncrement i m = do
i' <- unsafeGetElem c i
if i' /= m then return (Just (i,i')) else recurse
where
recurse = if i /= 0 then findIncrement (i-1) (m-1) else return Nothing
setAscending k i x | i == k = return ()
| otherwise = do
unsafeSetElem c i x
setAscending k (i+1) (x+1)
{-# INLINE setNext #-}
setPrev :: (MChoose c m) => c -> m Bool
setPrev c = do
n <- getPossible c
k <- getSize c
if k > 0
then do
k1' <- unsafeGetElem c (k-1)
findGap (k-1) k1' >>=
maybe (return False) (\(i,i') -> do
unsafeSetElem c i (i'-1)
setAscending k (i+1) (n-k+i+1)
return True
)
else
return False
where
findGap i i'
| i == 0 =
if i' == 0
then return $ Nothing
else return $ Just (0,i')
| otherwise = let j = i-1 in do
j' <- unsafeGetElem c j
if i' /= j'+1
then return $ Just (i,i')
else findGap j j'
setAscending k i x | i == k = return ()
| otherwise = do
unsafeSetElem c i x
setAscending k (i+1) (x+1)
{-# INLINE setPrev #-}
freeze :: (MChoose c m) => c -> m Choose
freeze c = unsafeFreeze =<< newCopyChoose c
{-# INLINE freeze #-}
thaw :: (MChoose c m) => Choose -> m c
thaw c = newCopyChoose =<< unsafeThaw c
{-# INLINE thaw #-}
instance MChoose (STChoose s) (ST s) where
getPossible = getPossibleSTChoose
{-# INLINE getPossible #-}
getSize = getSizeSTChoose
{-# INLINE getSize #-}
newChoose = newSTChoose
{-# INLINE newChoose #-}
newChoose_ = newSTChoose_
{-# INLINE newChoose_ #-}
unsafeGetElem = unsafeGetElemSTChoose
{-# INLINE unsafeGetElem #-}
unsafeSetElem = unsafeSetElemSTChoose
{-# INLINE unsafeSetElem #-}
getElems = getElemsSTChoose
{-# INLINE getElems #-}
setElems = setElemsSTChoose
{-# INLINE setElems #-}
unsafeFreeze = unsafeFreezeSTChoose
{-# INLINE unsafeFreeze #-}
unsafeThaw = unsafeThawSTChoose
{-# INLINE unsafeThaw #-}
instance MChoose IOChoose IO where
getPossible = getPossibleIOChoose
{-# INLINE getPossible #-}
getSize = getSizeIOChoose
{-# INLINE getSize #-}
newChoose = newIOChoose
{-# INLINE newChoose #-}
newChoose_ = newIOChoose_
{-# INLINE newChoose_ #-}
unsafeGetElem = unsafeGetElemIOChoose
{-# INLINE unsafeGetElem #-}
unsafeSetElem = unsafeSetElemIOChoose
{-# INLINE unsafeSetElem #-}
getElems = getElemsIOChoose
{-# INLINE getElems #-}
setElems = setElemsIOChoose
{-# INLINE setElems #-}
unsafeFreeze = unsafeFreezeIOChoose
{-# INLINE unsafeFreeze #-}
unsafeThaw = unsafeThawIOChoose
{-# INLINE unsafeThaw #-}