Ticket #3320: idfs.hs

{-# LANGUAGE RankNTypes #-}
{-# OPTIONS_GHC -O -fno-full-laziness #-}

import Control.Parallel            ( par )
import Control.Parallel.Strategies ( using, seqList, r0 )

newtype Search a = S { search :: forall b . (a -> Int -> [b]) -> Int -> [b] }

idfs :: Int -> Search a -> [a]
idfs n a = parconcat [ search a result d | d <- [0,n..] ]
 where result x d = if d<n then [x] else []

parconcat :: [[a]] -> [a]
parconcat []     = []
parconcat (x:xs) = ys `par` ((x `using` seqList r0) ++ ys)
 where ys = parconcat xs

ret :: a -> Search a
ret x = S ($x)

bind :: Search a -> (a -> Search b) -> Search b
bind (S a) k = S (\c -> a (\x -> search (k x) c))

zero :: Search a
zero = S (\_ _ -> [])

plus :: Search a -> Search a -> Search a
plus (S a) (S b) = S (\c d -> if d==0 then [] else a c (d-1) ++ b c (d-1))

anyof :: [a] -> Search a
anyof = foldr plus zero . map ret

pytriple :: Search (Int,Int,Int)
pytriple = anyof [1..]   `bind` \a ->
           anyof [a+1..] `bind` \b ->
           anyof [b+1..] `bind` \c ->
           if a*a + b*b == c*c then ret (a,b,c) else zero

main = print . length . take 500 . idfs 100 $ pytriple