Haskell checkout kata II

Main.hs

{- a simpler implementation -}

module Main (main) where

import Test.Hspec
import Test.Hspec.QuickCheck
import Data.List
import qualified Data.Map as M

main :: IO ()
main = hspec $ do

  describe "calcTotal" $ do
    let rules = [('A', 1, 50), ('A', 3, 130), ('B', 1, 30)]
    it "returns zero with no items" $ calcTotal rules [] `shouldBe` 0
    it "handles one item"           $ calcTotal rules ['A'] `shouldBe` 50
    it "handles a different item"   $ calcTotal rules ['B'] `shouldBe` 30
    it "handles two of an item"     $ calcTotal rules ['A', 'A'] `shouldBe` 100
    it "discounts three of an item" $ calcTotal rules ['A', 'A', 'A'] `shouldBe` 130
    it "discounts four of an item"  $ calcTotal rules ['A', 'A', 'A', 'A'] `shouldBe` 180

  describe "groupItems" $ do
    it "groups an empty collection" $ groupItems [] `shouldBe` []
    it "groups a singleton"         $ groupItems ['A'] `shouldBe` [('A', 1)]

    prop "group/ungroup == sort" $ \xs -> concat [replicate count sku | (sku, count) <- groupItems xs] `shouldBe` sort xs

type Sku = Char
type Price = Double
type PricingRule = (Sku, Int, Price)

calcTotal :: [PricingRule] -> [Sku] -> Price
calcTotal pricingRules items = sum $ map calcGroupPrice $ groupItems items
  where
  pricingRulesMap :: M.Map Sku (M.Map Int Price)
  pricingRulesMap = M.fromListWith (M.unionWith min)
    [ (sku, M.singleton count price)
    | (sku, count, price) <- pricingRules
    ]

  calcGroupPrice (sku, count0) = case M.lookup sku pricingRulesMap of
      Nothing -> error $ "never heard of sku " ++ show sku
      Just dealsMap -> calcDealPrice dealsMap count0

  calcDealPrice dealsMap = calcCountPrice
    where
    calcCountPrice count
      | count <= 0 = 0
      | otherwise = case M.lookupLE count dealsMap of
        Nothing -> error $ "no price for <= " ++ show count ++ " of something"
        (Just (dealCount, dealPrice)) -> dealPrice + calcCountPrice (count - dealCount)

groupItems :: [Sku] -> [(Sku, Int)]
groupItems skus = [(sku, length thisGroup) | thisGroup@(sku:_) <- group $ sort skus]