Haskell checkout kata

Main.hs

{-# LANGUAGE GeneralizedNewtypeDeriving #-}
{-# LANGUAGE LambdaCase #-}
{-# LANGUAGE RecordWildCards #-}

module Main (main) where

import Control.Applicative
import Control.Monad
import Control.Monad.RWS
import Control.Monad.Trans.Either
import Data.Void
import Test.Hspec
import qualified Data.Map as M

data PricingRule sku q p = PricingRule
  { sku      :: sku
  , quantity :: q
  , price    :: p
  }

main :: IO ()
main = hspec $ do
  let rules =
        [ PricingRule 'A' 1 50
        , PricingRule 'A' 3 130
        , PricingRule 'B' 1 30
        , PricingRule 'B' 2 45
        , PricingRule 'C' 1 20
        , PricingRule 'D' 1 15
        :: PricingRule Char Int Int]
  describe "api" $ do
    it "has a nice api" $ runCheckout rules (do
        scanItem 'A'
        scanItem 'D'
        scanItem 'B'
        t1 <- calcTotal
        scanItem 'A'
        scanItem 'B'
        scanItem 'A'
        t2 <- calcTotal
        scanItem 'E'
        t3 <- calcTotal
        return (t1, t2, t3)) `shouldBe` (TotalResult 95, TotalResult 190, NoPricing 'E' 0)

    it "returns zero on empty"      $ runCheckout rules calcTotal `shouldBe` TotalResult 0

    it "handles an unknown item"    $ runCheckout rules (do
      scanItem 'E'
      calcTotal) `shouldBe` NoPricing 'E' 0

    let rulesWithoutSingleE = PricingRule 'E' 2 30 : rules

    it "handles too few of an item"   $ runCheckout rulesWithoutSingleE (do
      scanItem 'E'
      calcTotal) `shouldBe` NoPricing 'E' 1
    it "copes with enough of an item" $ runCheckout rulesWithoutSingleE (do
      scanItem 'E'
      scanItem 'E'
      calcTotal) `shouldBe` TotalResult 30

  let skus `shouldTotal` expectedPrice
        = it worksAsExpected $ runCheckout rules (mapM_ scanItem skus >> calcTotal) `shouldBe` TotalResult expectedPrice
        where
        worksAsExpected = skus ++ " should total to " ++ show expectedPrice

  describe "totals" $ do
    ""       `shouldTotal` 0
    "A"      `shouldTotal` 50
    "AB"     `shouldTotal` 80
    "CDBA"   `shouldTotal` 115
    "AA"     `shouldTotal` 100
    "AAA"    `shouldTotal` 130
    "AAAA"   `shouldTotal` 180
    "AAAAA"  `shouldTotal` 230
    "AAAAAA" `shouldTotal` 260
    "AAAB"   `shouldTotal` 160
    "AAABB"  `shouldTotal` 175
    "AAABBD" `shouldTotal` 190
    "DABABA" `shouldTotal` 190

type PricingRuleMap sku q p = M.Map sku (M.Map q p)

type ItemTally sku q = M.Map sku q

newtype CheckoutM sku q p a = CheckoutM (RWS (PricingRuleMap sku q p) () (ItemTally sku q) a)
  deriving (Functor, Applicative, Monad)

runCheckout :: (Ord q, Ord sku) => [PricingRule sku q p] -> CheckoutM sku q p a -> a
runCheckout prs (CheckoutM go) = case runRWS go prsMap M.empty of (a,_,_) -> a
  where prsMap = M.fromListWith M.union [(sku, M.singleton quantity price) | PricingRule{..} <- prs]

scanItem :: (Num q, Ord sku) => sku -> CheckoutM sku q p ()
scanItem c = CheckoutM $ modify $ M.insertWith (+) c 1

data TotalResult sku q p
  = TotalResult p
  | NoPricing sku q
  deriving (Eq, Show)

eitherVoid :: Either a Void -> a
eitherVoid = either id absurd

calcTotal :: (Ord q, Num q, Ord sku, Num p) => CheckoutM sku q p (TotalResult sku q p)
calcTotal = run <$> CheckoutM ask <*> CheckoutM get
  where
  run = runCalcTotalM $ forever $ do
    (sku, count)           <- takeNextItems
    (dealCount, dealPrice) <- lookupPricing sku count
    addToTotal dealPrice
    let leftover = count - dealCount
    when (leftover > 0) $ replaceItems sku leftover

newtype CalcTotalM sku q p a = CalcTotalM (EitherT (Sum p -> TotalResult sku q p) (RWS (PricingRuleMap sku q p) (Sum p) (ItemTally sku q)) a)
  deriving (Functor, Applicative, Monad)

runCalcTotalM :: CalcTotalM sku q p Void -> PricingRuleMap sku q p -> ItemTally sku q -> TotalResult sku q p
runCalcTotalM (CalcTotalM erwsa) prm ity = case runRWS (runEitherT erwsa) prm ity of (mf, _, t) -> eitherVoid mf t

addToTotal :: Num p => p -> CalcTotalM sku q p ()
addToTotal = CalcTotalM . tell . Sum

getItemTally :: Num p => CalcTotalM sku q p (ItemTally sku q)
getItemTally = CalcTotalM get

putItemTally :: Num p => ItemTally sku q -> CalcTotalM sku q p ()
putItemTally = CalcTotalM . put

replaceItems :: (Ord sku, Num p) => sku -> q -> CalcTotalM sku q p ()
replaceItems sku count = CalcTotalM $ modify $ M.insert sku count

withResult :: Sum p -> TotalResult sku q p
withResult = TotalResult . getSum

withNoPricing :: sku -> q -> Sum p' -> TotalResult sku q p
withNoPricing sku count _ = NoPricing sku count

askPricingRules :: Num p => CalcTotalM sku q p (PricingRuleMap sku q p)
askPricingRules = CalcTotalM ask

maybeExit :: Num p => (Sum p -> TotalResult sku q p) -> Maybe a -> CalcTotalM sku q p a
maybeExit f Nothing  = exitCalc f
maybeExit _ (Just a) = return a

exitCalc :: Num p => (Sum p -> TotalResult sku q p) -> CalcTotalM sku q p a
exitCalc f = CalcTotalM $ left f

takeNextItems :: Num p => CalcTotalM sku q p (sku, q)
takeNextItems = do
  itemTally <- getItemTally
  ((sku, count), remainingItems) <- maybeExit withResult $ M.minViewWithKey itemTally
  putItemTally remainingItems
  return (sku, count)

lookupItemPricing :: (Ord sku, Num q, Num p) => sku -> CalcTotalM sku q p (M.Map q p)
lookupItemPricing sku = do
  pricingRules <- askPricingRules
  maybeExit (withNoPricing sku 0) $ M.lookup sku pricingRules

lookupPricing :: (Ord sku, Ord q, Num q, Num p) => sku -> q -> CalcTotalM sku q p (q, p)
lookupPricing sku count = do
  itemPricing <- lookupItemPricing sku
  maybeExit (withNoPricing sku count) $ M.lookupLE count itemPricing