Code from Chapter 6 of Pearls of Functional Algorithm Design

pearls-ch6.hs

#!/usr/bin/env stack
-- stack --install-ghc runghc --package=criterion

import Criterion.Main
import Data.Char (intToDigit)
import Data.List (intercalate)

type Expression = [Term]
type Term = [Factor]
type Factor = [Digit]
type Digit = Int

valExpr :: Expression -> Int
valExpr = sum . map valTerm

valTerm :: Term -> Int
valTerm = product . map valFact

valFact :: Factor -> Int
valFact = foldl1 (\n d -> 10 * n + d)

showExpr :: Expression -> String
showExpr = intercalate " + " . map showTerm

showTerm :: Term -> String
showTerm = intercalate "×" . map showFact

showFact :: Factor -> String
showFact = map intToDigit

good :: Int -> Int -> Bool
good c v = v == c

ok :: Int -> Int -> Bool
ok c v = v <= c

expressions :: [Digit] -> [Expression]
expressions = foldr extend []

extend :: Digit -> [Expression] -> [Expression]
extend x [] = [[[[x]]]]
extend x es = concatMap (glue x) es

glue :: Digit -> Expression -> [Expression]
glue x ((xs : xss) : xsss) =
    [ ((x : xs) : xss) : xsss
    , ([x] : xs : xss) : xsss
    , [[x]] : (xs : xss) : xsss
    ]

digits0 :: [Int]
digits0 = [1..9]

target0 :: Int
target0 = 100

solutions :: Int -> [Digit] -> [Expression]
solutions c ds = filter (good c . valExpr) $ expressions ds

showSolution :: Int -> Expression -> String
showSolution c x = show c ++ " = " ++ showExpr x

printSolutions :: Int -> [Expression] -> IO ()
printSolutions c xs = mapM_ (putStrLn . showSolution c) xs

computeAndPrintSolutions :: Int -> [Digit] -> IO ()
computeAndPrintSolutions c ds = printSolutions c $ solutions c ds

piDigits :: [Int]
piDigits = [3, 1, 4, 1, 5, 9, 2, 6, 5, 3, 5, 8, 9, 7]

target1 :: Int
target1 = 1000

cross :: (a -> b, c -> d) -> (a, c) -> (b, d)
cross (f, g) (x, y) = (f x, g y)

modify :: Digit -> (Int, Int, Int, Int) -> [(Int, Int, Int, Int)]
modify x (k, f, t, e) =
    [(10*k, k*x+f, t, e), (10, x, f*t, e), (10, x, 1, f*t+e)]

good' :: Int -> (Int, Int, Int, Int) -> Bool
good' c (k, f, t, e) = f*t + e == c

ok' :: Int -> (Int, Int, Int, Int) -> Bool
ok' c (k, f, t, e) = f*t + e <= c

solutions' :: Int -> [Digit] -> [Expression]
solutions' c = map fst . filter (good' c . snd) . foldr (expand c) []

expand :: Int -> Digit -> [(Expression, (Int, Int, Int, Int))] -> [(Expression, (Int, Int, Int, Int))]
--expand c x = filter (ok' c . snd) . uncurry zip . cross (extend x, modify x) . unzip
expand c x [] = [([[[x]]], (10, x, 1, 0))]
expand c x evs = concat (map (filter (ok' c . snd) . glue' x) evs)

glue' :: Digit -> (Expression, (Int, Int, Int, Int)) -> [(Expression, (Int, Int, Int, Int))]
glue' x ((xs : xss) : xsss, (k, f, t, e)) =
    [ (((x : xs) : xss) : xsss, (10*k, k*x+f, t, e))
    , (([x] : xs : xss) : xsss, (10, x, f*t, e))
    , ([[x]] : (xs : xss) : xsss, (10, x, 1, f*t+e))
    ]

main = defaultMain
    [ bgroup "solutions"
        [ bench "initial version, target 100, digits 1..9" $ nf (solutions 100) digits0
        , bench "initial version, target 1000, pi digits" $ nf (solutions 1000) piDigits
        , bench "improved version, target 1000, pi digits" $ nf (solutions' 1000) piDigits
        ]
    ]