Dbs.hs

{-# LANGUAGE FlexibleInstances #-}
{-# LANGUAGE RankNTypes        #-}
{-# LANGUAGE TypeFamilies      #-}

module Dbs where

import           Control.Monad.Morph
import           Control.Monad.ST         (RealWorld, ST, runST, stToIO)
import           Control.Monad.Trans.Free
import           Data.Array.ST
import           Lens.Family
import           Pipes                    hiding (yield)
import qualified Pipes                    as P
import           Pipes.Group              (chunksOf, intercalates)
import qualified Pipes.Prelude            as P

-- calculate de bruijn sequence

class Monad m => STW m where
    type Yield m
    type RawState m
    yield    :: Yield m -> m ()
    readArr  :: (Ix i) => STArray (RawState m) i e -> i -> m e
    writeArr :: (Ix i) => STArray (RawState m) i e -> i -> e -> m ()
    newArr   :: (Ix i) => (i,i) -> e -> m (STArray (RawState m) i e)

instance STW (Producer w (ST s)) where
  type Yield (Producer w (ST s)) = w
  type RawState (Producer w (ST s)) = s
  yield          = P.yield
  readArr a i    = lift (readArray a i)
  writeArr a i e = lift (writeArray a i e)
  newArr ix e    = lift (newArray ix e)

db :: (STW m, Yield m ~ Integer) => Integer -> Integer -> STArray (RawState m) Integer Integer -> Integer -> Integer -> m ()
db k n a t p
    | t <= n = do
        readArr a (t-p) >>= writeArr a t
        db k n a (t+1) p
        start <- readArr a (t-p)
        flip mapM_ [start+1..k-1] $ \j -> do
            writeArr a t j
            db k n a (t+1) t
    | n `mod` p == 0 = flip mapM_ [1..p] $ \j -> do
        v <- readArr a j
        yield v
    | otherwise = return ()

debruijnM :: (STW m, Yield m ~ Integer) => Integer -> Integer -> m ()
debruijnM k n = do
    a <- newArr (0::Integer,k*n-1) 0
    db k n a 1 1

vars :: [Char]
vars = " @_" ++ ['0'..'9'] ++ ['a'..'z'] ++ ['A'..'Z']

testP
  :: (Integer -> Integer -> Producer Integer (ST RealWorld) ())
     -> Integer -> Integer -> IO ()
testP f k n =
  runEffect $ for charStream (lift . putChar)
  where
    charStream = intercalates (P.yield '\n') chunkedIO
    chunkedIO = hoist stToIO producerC ^. chunksOf (fromIntegral n)
    producerC = producerI >-> P.map (\x -> vars !! fromInteger x)
    producerI = f k n