Dragon Jousting Riddle in Ersatz

DragonJoustingRiddle.hs

#!/usr/bin/env stack
{- stack
     --resolver lts-15.15
     script
     --package ersatz
     --package containers
-}

{-# LANGUAGE RecordWildCards #-}
{-# LANGUAGE TypeApplications #-}
{-# LANGUAGE TypeFamilies #-}

module Main (main) where

-- https://youtu.be/KSkw7hKN_Xg

import Data.Foldable
import qualified Data.Map.Strict as M
import Ersatz

data Team = Elves | Goblins | Treefolk
  deriving (Eq, Ord, Enum, Bounded, Show, Read)

type BitMatches = M.Map (Team, Team) BitMatch

data Match = NoMatch
           | Match Integer Integer
  deriving (Eq, Ord, Show)

data BitMatch = BitMatch{ matchOccurred :: Bit
                        , matchScored   :: Bits
                        , matchReceived :: Bits
                        }
  deriving (Show)

instance Codec BitMatch where
  type Decoded BitMatch = Match

  decode s BitMatch{..} = do
    occurred <- decode s matchOccurred
    if occurred
      then Match <$> decode s matchScored <*> decode s matchReceived
      else pure NoMatch

  encode NoMatch = BitMatch (encode false) (encode 0) (encode 0)
  encode (Match scored received) =
    BitMatch (encode true) (encode scored) (encode received)

main :: IO ()
main = do
  ~(Satisfied, Just ms) <- solveWith anyminisat bitMatches
  traverse_ print (M.toList ms)

bitMatches :: MonadSAT s m => m BitMatches
bitMatches = do
  msUnique <- M.fromList
          <$> traverse (\team -> (,) team <$> bitMatch) teamCombinations

  let ms = msUnique `M.union` (M.fromList . fmap flipMatch . M.toList) msUnique

  assert . Ersatz.and $
    [ (sumBit  . fmap matchOccurred . filterMatches Elves)    ms >=? 1
    , (sumBit  . fmap matchOccurred . filterMatches Goblins)  ms >=? 1
    , (sumBit  . fmap matchOccurred . filterMatches Treefolk) ms >=? 1

    , (sumBit  . fmap matchOccurred . filterMatches Elves)    ms /== 1
    , (sumBit  . fmap matchWon      . filterMatches Elves)    ms /== 0
    , (sumBits . fmap matchScored   . filterMatches Elves)    ms /== 6

    , (sumBit  . fmap matchOccurred . filterMatches Goblins)  ms /== 2
    , (sumBit  . fmap matchTied     . filterMatches Goblins)  ms /== 1
    , (sumBits . fmap matchScored   . filterMatches Goblins)  ms /== 0
    , (sumBits . fmap matchReceived . filterMatches Goblins)  ms /== 3

    , (sumBit  . fmap matchTied     . filterMatches Treefolk) ms /== 0
    , (sumBits . fmap matchScored   . filterMatches Treefolk) ms /== 0
    , (sumBits . fmap matchReceived . filterMatches Treefolk) ms /== 1
    ]

  pure msUnique

  where
    -- Unique combinations of teams, teamA < teamB.
    teamCombinations = do
      teamA <- [minBound :: Team ..]
      teamB <- drop 1 [teamA ..]
      pure (teamA, teamB)

    flipMatch ((teamA, teamB), BitMatch{..}) =
      ((teamB, teamA), flippedMatch)
      where flippedMatch = BitMatch{ matchOccurred = matchOccurred
                                   , matchScored   = matchReceived
                                   , matchReceived = matchScored
                                   }

    filterMatches team =
      M.filterWithKey (\(teamA, _teamB) _match -> teamA == team)

bitMatch :: MonadSAT s m => m BitMatch
bitMatch = do
  occurred <- exists
  scored   <- bits <$> exists @Bit4  -- Always <= 5.
  received <- bits <$> exists @Bit4
  assert $ choose (Ersatz.and [ scored === 0, received === 0 ])
                  (Ersatz.and [ scored >=? 0, received >=? 0, scored + received <=? 5 ])
                  occurred
  pure (BitMatch occurred scored received)

matchWon, matchTied :: BitMatch -> Bit
matchWon  BitMatch{..} = matchOccurred Ersatz.&& (matchScored >?  matchReceived)
matchTied BitMatch{..} = matchOccurred Ersatz.&& (matchScored === matchReceived)