implementing sequencePipes

p8.hs

import System.Environment
import qualified Data.Array as A
import Data.Array.IO
import qualified Data.Array.MArray as M
import Data.Array.Base  ( UArray, unsafeFreezeSTUArray )
import Data.Array.Unboxed
import qualified Data.Array.Unboxed as U
import Control.Monad
import Debug.Trace
import Pipes
import qualified Pipes.Prelude as P
import System.IO
import Data.IORef

swap a i j = do
  t <- readArray a i
  readArray a j >>= writeArray a i
  writeArray a j t

reverseA a i j
  | i < j = do swap a i j
               reverseA a (i+1) (j-1)
  | otherwise = return ()

-- next permutation
nextP :: Int -> IOUArray Int Int -> IO Bool
nextP n p = do
  -- find largest k s.t. p[k] < p[k+1]
  fp <- M.freeze p :: IO (UArray Int Int)
  let ks = [ k | k <- [n-1,n-2..1], fp!k < fp!(k+1) ]
  case ks of
    []    -> return False
    (k:_) -> do let j = head [ j | j <- [n,n-1..k+1], fp!j > fp!k ]
                swap p k j
                reverseA p (k+1) n
                return True

allPerms :: Int -> Producer [Int] IO ()
allPerms n = do
  p <- lift $ (newListArray (1,n) [1..n] :: IO (IOUArray Int Int))
  lift (getElems p) >>= yield
  let loop = do b <- lift $ nextP n p
                if b
                  then do { lift (getElems p) >>= yield; loop }
                  else return ()
  loop

sgen :: Int -> Producer Int IO ()
sgen n = do
  r <- lift $ newIORef 1
  let loop = do v <- lift $ readIORef r
                if v <= n
                  then do { yield v; lift $ modifyIORef r (+1); loop }
                  else return ()
  loop

sequencePipes :: Monad m => [ Producer a m () ] -> Producer [a] m ()
sequencePipes [] = yield []
sequencePipes (m:ms) = do
  for m $ \x -> do
    for (sequencePipes ms) $ \xs -> yield (x:xs)

testN = runEffect $ for (sequencePipes [sgen 3, sgen 2, sgen 4]) $ lift . print

ok [] = True
ok (a:as) = all (< a) as && ok as

test1 n = do
  putStrLn $ "sgen m, where m = " ++ show m
  P.last (sgen m) >>= print
    where m = 2^(n :: Int)

test2 n = do
  putStrLn $ "allPerms n, where n = " ++ show n
  P.last (allPerms n) >>= print

test3a n = do
  putStrLn $ "sequencePipes $ replicate 2 $ sgen m, where m = " ++ show m
  P.last (sequencePipes $ replicate 2 $ sgen m) >>= print
  where m = 2^(n::Int)

test3b n = do
  putStrLn $ "for (sgen m) ... for (sgen m) ..., where m = " ++ show m
  P.last (for (sgen m) $ \x1 -> for (sgen m) $ \x2 -> yield [x1,x2]) >>= print
  where m = 2^(n::Int)

test4 n = do
  putStrLn $ "sequencePipes $ replicate 3 $ allPerms n, where n = " ++ show n
  P.last pipes >>= print
  where pipes = sequencePipes $ replicate 3 $ allPerms n

test5 n = do
  putStrLn $ "for (allPerms n) ... for (allPerms n) ..., where n = " ++ show n
  let pipe = for (allPerms n) $ \p1 -> for (allPerms n) $ \p2 -> yield [p1,p2]
  P.last pipe >>= print

test6 n = do
  putStrLn $ "sequencePipes $ replicate 2 $ allPerms n, where n = " ++ show n
  P.last pipes >>= print
  where pipes = sequencePipes $ replicate 2 $ allPerms n

main = do
  (testno:n:_) <- fmap (map read) getArgs
  case testno of
    1 -> test1 n
    2 -> test2 n
    4 -> test4 n
    5 -> test5 n
    6 -> test6 n
    7 -> test3a n
    8 -> test3b n
    _ -> error "bad test number"