shapes.hs

{- cabal:
build-depends: base, bluefin
-}
{-# LANGUAGE GHC2021 #-}

import Bluefin.Eff (Eff, runPureEff, (:>))
import Bluefin.Stream (Stream, forEach, yield, yieldToList)
import Control.Monad (when)
import Data.Foldable (for_)

replicateList :: Int -> [a] -> [[a]]
replicateList n l = fst (runPureEff (yieldToList (replicateListEff n l)))

replicateListEff :: (e1 :> es) => Int -> [a] -> Stream [a] e1 -> Eff es ()
replicateListEff 0 _ y = yield y []
replicateListEff n l y =
  for_ l $ \x -> do
    forEach (\y' -> replicateListEff (n - 1) l y') $ \sub -> do
      yield y (x : sub)

-- Using this one leaks space
shapesEffLeak ::
  (e1 :> es) => Integer -> Integer -> Stream [Integer] e1 -> Eff es ()
shapesEffLeak bias p y = do
  for_ [0 .. p] $ \m -> do
    for_ (replicateList (fromInteger m + 2) [1 .. p]) $ \s -> do
      when (p == sum (zipWith (*) (map (bias +) s) (tail s))) $ do
        yield y s

-- Using this one runs in constant space
shapesEff ::
  (e1 :> es) => Integer -> Integer -> Stream [Integer] e1 -> Eff es ()
shapesEff bias p y = do
  for_ [0 .. p] $ \m -> do
    forEach (replicateListEff (fromInteger m + 2) [1 .. p]) $ \s -> do
      when (p == sum (zipWith (*) (map (bias +) s) (tail s))) $ do
        yield y s

shapes :: Integer -> Integer -> [[Integer]]
shapes bias p = fst (runPureEff (yieldToList (shapesEff bias p)))

main :: IO ()
main = mapM_ (print . length . shapes 0) [1 .. 10]