Vigenere cipher exercise in Haskell

Question

This is my implementation using the dreadful !!:

import Data.Char (chr, ord, toUpper)


-- A bit of self documentation help
type Key = String
type Msg = String


key :: Key
key = "TSTING"
msg :: Msg
msg = "I'm not even mad... This is impressive!"


-- | Checks if character is valid for encoding
isValid :: Char -> Bool
isValid c = let cUp = toUpper c :: Char
             in 'A' <= cUp && cUp <= 'Z'


-- | Given 'key' & 'msg' generate a list of [Maybe Int] indices
-- to map 'msg' from 'key', skipping invalid characters
toIdx :: Key -> Msg -> [Maybe Int]
toIdx k m = map (flip mod keyN <$>) $ toIdx_ 0 m
  where keyN = length k :: Int
        toIdx_ :: Int -> Msg -> [Maybe Int]
        toIdx_ _ "" = []
        toIdx_ acc (c:cs)
          | isValid c = Just acc : toIdx_ (acc + 1) cs
          | otherwise = Nothing : toIdx_ acc cs


-- | Given 'key' & 'msg' generate a list of numbers representing
-- the amount to shift 'msg' characters based on 'key'
toShifts :: Key -> Msg -> [Int]
toShifts k m = map toKey (toIdx k m)
  where kUp = map toUpper k :: Key
        toKey :: Maybe Int -> Int
        toKey Nothing  = 0
        toKey (Just x) = ord (kUp!!x) - ord 'A'


-- | Given 'by' & 'c', shift the Char 'c' by amount 'by'. 'by' can be both
-- positive & negative as well as 0.
shift :: Int -> Char -> Char
shift by c
  | isValid c && c >= 'a' = shift_ $ ord 'a'
  | isValid c && c >= 'A' = shift_ $ ord 'A'
  | otherwise = c
  where cONorm    = ord (toUpper c) - ord 'A' :: Int
        azN       = ord 'Z' - ord 'A' :: Int
        shift_ :: Int -> Char
        shift_ aO = chr $ aO + mod (by + cONorm) azN


-- Encode & decode a message using the given key.
vigenere, unVigenere :: Key -> Msg -> Msg
vigenere   k m = zipWith shift (toShifts k m) m
unVigenere k m = zipWith shift (map negate $ toShifts k m) m

I found that the most "annoying" thing when coming from background such as Python is to be able to keep track of things, for example when figuring out how to convert valid characters into usable positions to be then mapped with the key. That thing took me half a day to figure out!

How would you do it? Or is there some "standard" way of working with these sort of things? I'm referring particularly to toIdx & toIdx_, had to use toIdx_ to accumulate the indices with acc in this list of Maybe Int in order to correctly map valid Chars with key.

Of course, I wouldn't have it any other way but have an algorithm which takes any String input and creates an encoded String with preserved upper/lower-case and non-valid (out of ASCII alphabet) Chars.

Answer 1

You're rather close. The issue lies within zipWith. zipWith will always consume the first elements of both lists. That's why you need (toShifts k m) to begin with.

However, what if we get rid of zipWith for a moment and use pattern matching in vigenere?

vigenere' :: Key -> Msg -> Msg
vigenere' ks     (' ':ms) = ' '        : vigenere ks ms
vigenere' (k:ks) (m  :ms) = shiftC k m : vigenere ks ms
vigenere' _      []       = []

where shiftC :: Char -> Char -> Char is an appropriate function (left as an exercise). Actually, that's it. That's all that is necessary for vigenere, apart from shiftC. Well, I'm lying: the key will run out at some point. That's why you use cycle :: [a] -> [a]:

vigenere :: Key -> Msg -> Msg
vigenere ks ms = vigenere' (cycle ks) ms

cycle turns a regular list in an infinite one, e.g. cycle [1,2,3] == [1,2,3,1,2,3,1,2,3,…] (exercise: try to implement cycle). Therefore, it turns your key into an infinite one.

You can implement unVigenere' the same way, only that you need unShiftC :: Char -> Char -> Char. A pair of functions, namely toInt :: Char -> Int and fromInt :: Int -> Char will come in handy for that (left as an exercise).

Heck, we can even implement both functions the same way and use isValid instead of pattern matching for unsupported characters:

cryptZip' :: (Char -> Char -> Char) -> Key -> Msg -> Msg
cryptZip' _ _      []       = []
cryptZip' f (k:ks) (m  :ms) =
  | isValid m = f k m : cryptZip' f ks     ms
  | otherwise = m     : cryptZip' f (k:ks) ms

-- | Combines the key and the message with the given function.
-- Invalid characters are left as-is. The function shall return
-- only valid characters.
cryptZip :: (Char -> Char -> Char) -> Key -> Msg -> Msg
cryptZip f ks ms = cryptZip' f (cycle ks) ms

Now, vigenere and unVigenere are extremely simple:

vigenere   = cryptZip shiftC
unVigenere = cryptZip unShiftC

And we're done. By the way, you can use (Char -> Char -> Maybe Char) to get rid of isValid. But that, again, is left as an exercise.

Further remarks

The type annotations in your where bindings are not necessary. I would get rid of them, since they may get out of sync with your top-level function (but you will at least get a type-error).

cryptZip can be written as (a -> a -> Maybe a) -> [a] -> [a] -> [a], if you're up for a challenge.

Your overall complexity stemmed from zipWith. It's a great function, but only if it's used for the right job. Not everything is a nail, just because you have a hammer at hand. Instead of trying to get everything else in place to use a specific function, ask yourself what you want to do, and whether there's an easy way. Unfortunately, there is no cryptZip-like function in the standard library (unless you count stateful accumulators such as mapAccumL).

Answer 2

To get rid of !! here, you can use it earlier and earlier until you never even generate an Int.

-- | Given 'key' & 'msg' generate a list of [Maybe Int] indices
-- to map 'msg' from 'key', skipping invalid characters
toIdx :: Key -> Msg -> [Maybe Char]
toIdx k m = toIdx_ (cycle $ map toUpper k) m
  where toIdx_ :: Key -> Msg -> [Maybe Char]
        toIdx_ _ "" = []
        toIdx_ key@(k:ey) (c:cs)
          | isValid c = Just k : toIdx_ ey cs
          | otherwise = Nothing : toIdx_ key cs


-- | Given 'key' & 'msg' generate a list of numbers representing
-- the amount to shift 'msg' characters based on 'key'
toShifts :: Key -> Msg -> [Int]
toShifts k m = map toKey (toIdx k m)
  where toKey :: Maybe Char -> Int
        toKey Nothing  = 0
        toKey (Just x) = ord x - ord 'A'

Of course, there is hardly a need to separate all these steps.

base :: Char -> Maybe Char
base c
  | 'a' <= c && c <= 'z' = Just 'a'
  | 'A' <= c && c <= 'Z' = Just 'A'
  | otherwise = Nothing

vigenere, unVigenere :: String -> String -> String
[vigenere, unVigenere] = (`map` [(+), (-)]) $ \direction k ->
  (.) snd $ (`mapAccumL` cycle k) $ \key@(k:ey) c -> case base c of
    Nothing -> (key, c)
    Just a -> (,) ey $ chr $ ord a +
      mod ((ord c - ord a) `direction` (ord (toUpper k) - ord 'A')) (ord 'Z' - ord 'A')