"Einstein's Riddle" in Ersatz

EinsteinsRiddle.hs

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

{-# LANGUAGE DeriveGeneric #-}
{-# LANGUAGE GeneralizedNewtypeDeriving #-}
{-# LANGUAGE RecordWildCards #-}
{-# LANGUAGE TypeFamilies #-}

module Main (main) where

-- https://youtu.be/1rDVz_Fb6HQ

import Prelude hiding ((&&), and, any, all, true)

import Control.Monad
import Ersatz
import GHC.Generics (Generic)

data House = House { hNationality :: Nationality
                   , hWallColor   :: WallColor
                   , hCigar       :: Cigar
                   , hBeverage    :: Beverage
                   , hAnimal      :: Animal
                   }
  deriving (Eq, Ord, Show, Read)

data BitHouse = BitHouse { bhNationality :: BitEnum Nationality
                         , bhWallColor   :: BitEnum WallColor
                         , bhCigar       :: BitEnum Cigar
                         , bhBeverage    :: BitEnum Beverage
                         , bhAnimal      :: BitEnum Animal
                         }
  deriving (Show, Generic)

instance Codec BitHouse where
  type Decoded BitHouse = House
  decode s BitHouse{..} =
    House <$> decode s bhNationality <*> decode s bhWallColor
          <*> decode s bhCigar <*> decode s bhBeverage <*> decode s bhAnimal
  encode House{..} =
    BitHouse (encode hNationality) (encode hWallColor) (encode hCigar)
             (encode hBeverage) (encode hAnimal)

instance Equatable BitHouse
instance Orderable BitHouse
instance Variable BitHouse

data Nationality = Dane | Brit | German | Norwegian | Swede
  deriving (Eq, Ord, Bounded, Enum, Show, Read)

data WallColor = Blue | Green | Red | White | Yellow
  deriving (Eq, Ord, Bounded, Enum, Show, Read)

data Cigar = Blends | BlueMaster | Dunhill | PallMall | Prince
  deriving (Eq, Ord, Bounded, Enum, Show, Read)

data Beverage = Coffee | Milk | RootBeer | Tea | Water
  deriving (Eq, Ord, Bounded, Enum, Show, Read)

data Animal = Bird | Cat | Dog | Fish | Horse
  deriving (Eq, Ord, Bounded, Enum, Show, Read)

newtype BitEnum t = BitEnum Bits
  deriving (Show, Equatable, Orderable)

instance (Enum t) => Codec (BitEnum t) where
  type Decoded (BitEnum t) = t
  decode s (BitEnum b) = toEnum . fromInteger <$> decode s b
  encode = BitEnum . encode . toInteger . fromEnum

instance (Bounded t, Enum t) => Variable (BitEnum t) where
  literally m = do
    let beMaxBound = encode maxBound
        numBits = case beMaxBound of BitEnum (Bits bs) -> length bs
    be <- BitEnum . Bits <$> replicateM numBits (literally m)
    assert (be <=? beMaxBound)
    pure be

main :: IO ()
main = do
  ~(Satisfied, Just houses) <- solveWith anyminisat problem
  mapM_ print houses

problem :: MonadSAT s m => m [BitHouse]
problem = do
  houses <- replicateM 5 exists

  assert . and $
    [ unique (map bhNationality houses)
    , unique (map bhWallColor   houses)
    , unique (map bhCigar       houses)
    , unique (map bhBeverage    houses)
    , unique (map bhAnimal      houses)

    , any (\h        -> nationality h  Brit       && wallColor h  Red)      houses
    , any (\h        -> nationality h  Swede      && animal    h  Dog)      houses
    , any (\h        -> nationality h  Dane       && beverage  h  Tea)      houses
    , any (\(hl, hr) -> wallColor   hl Green      && wallColor hr White)    (neighborsLR houses)
    , any (\h        -> wallColor   h  Green      && beverage  h  Coffee)   houses
    , any (\h        -> cigar       h  PallMall   && animal    h  Bird)     houses
    , any (\h        -> wallColor   h  Yellow     && cigar     h  Dunhill)  houses
    , beverage    (houses !! 2) Milk
    , nationality (houses !! 0) Norwegian
    , any (\(h, hn)  -> cigar       h  Blends     && animal    hn Cat)      (neighbors houses)
    , any (\(h, hn)  -> animal      h  Horse      && cigar     hn Dunhill)  (neighbors houses)
    , any (\h        -> cigar       h  BlueMaster && beverage  h  RootBeer) houses
    , any (\h        -> nationality h  German     && cigar     h  Prince)   houses
    , any (\(h, hn)  -> nationality h  Norwegian  && wallColor hn Blue)     (neighbors houses)
    , any (\(h, hn)  -> cigar       h  Blends     && beverage  hn Water)    (neighbors houses)
    ]

  pure houses

  where
    nationality bh x = bhNationality bh === encode x
    wallColor   bh x = bhWallColor   bh === encode x
    cigar       bh x = bhCigar       bh === encode x
    beverage    bh x = bhBeverage    bh === encode x
    animal      bh x = bhAnimal      bh === encode x

unique :: Equatable a => [a] -> Bit
unique (x:xs) = all (x /==) xs && unique xs
unique [] = true

neighborsLR, neighbors :: [a] -> [(a, a)]
neighborsLR xs = zipWith (,) xs (drop 1 xs)
neighbors xs = do
  (l, r) <- neighborsLR xs
  [(l, r), (r, l)]