Haskell Quiz/The Solitaire Cipher/Solution Dolio

From HaskellWiki
Jump to navigation Jump to search


module Main where
import Data.List
import Data.Char
import Data.Ix
import Control.Monad
import Control.Monad.State
import System
import System.Random
import Data.Ord (comparing)

data Card = A | B | C Int deriving (Eq, Show)
type Deck = [Card]
type Cipher a = State Deck a

unkeyed = map C [1..52] ++ [A,B]

isJoker c = c == A || c == B

value (C i) = i
value _     = 53

-- en/decodes upper case characters into a 0 - 25 Int representation
-- (for easier arithmetic than 1 - 26)
decode n = chr (n + 65)
encode n = ord n - 65

-- Shuffles a given deck using n as the seed for the random generator
shuffle n = map snd . sortBy (comparing fst) . zip (randoms $ mkStdGen n :: [Int])

-- Scrubs/pads a string to turn it into the format expected by the cipher
scrub = join . intersperse " " . unfoldr split . pad . filter isAlpha . map toUpper
 where
 pad l = l ++ replicate ((5 - length l) `mod` 5) 'X'
 split [] = Nothing
 split l  = Just $ splitAt 5 l
  
-- Moves element e in l forward n places using the appropriate rules
push n e l = t ++ [e] ++ b
 where
 (Just i) = elemIndex e l
 l' = delete e l
 (t,b) = splitAt (n + i `mod` length l') l'

-- Performs the triple cut
tcut l = bottom ++ middle ++ top
 where
 [i,j] = findIndices isJoker l
 (top, m) = splitAt i l
 (middle, bottom) = splitAt (1 + j - i) m

-- Performs the counting cut
ccut l = init bottom ++ top ++ [last bottom]
 where
 n = value (last l)
 (top, bottom) = splitAt n l

-- Extracts a code from a given deck according to the appropriate rules.
-- Returns Nothing in the event that a joker is picked
extract l@(h:_) = if isJoker c then Nothing else Just (value c)
 where
 n = value h
 c = l !! n

-- Gets the next code in the key stream
getCode = do modify (ccut . tcut . push 2 B . push 1 A)
             deck <- get
             maybe getCode return (extract deck)

-- Uses the function f and initial deck d to en/decrypt a message
crypt f d = map decode . flip evalState d . mapM (cipher . encode)
 where
 cipher a 
    | inRange (0,25) a = getCode >>= return . flip mod 26 . f a
    | otherwise        = return a

decrypt = crypt (-)
encrypt = crypt (+)

crypto f = unlines . map (f . scrub) . lines

main = do (o:l) <- getArgs
          let deck = if null l then unkeyed else shuffle (read (head l)) unkeyed
          case o of
            "d" -> interact (crypto $ decrypt deck)
            "e" -> interact (crypto $ encrypt deck)
            _   -> putStrLn "Unrecognized option."