12
\$\begingroup\$

I wrote a Vigenere cipher in Haskell, but I feel like I way overdid it. This is the second version of the program that allows custom cipher tables, removing the limit of just having letters from A to Z, and nothing more.

I'd like to get some tips on what I could've done better, as this is the first program I've written in Haskell that is bigger than just a few functions.

module Vigenere where

import qualified Data.Map  as M
import qualified Data.List as L
import Data.Char

type Crypt a = a -> a -> CipherEnv -> Maybe a
type CharSet = [Char] -- alias for string, used to distinguish from text input
type CipherTable = M.Map Char CharSet

data CipherEnv = Cipher { cipherChars :: CharSet
                        , cipherTable :: CipherTable
                        } deriving (Show)

genRow :: Char -> CharSet -> Maybe CharSet
genRow char charset = let index = L.elemIndex char charset
                          stream = cycle charset
                          count  = length charset
                          nthSet = drop <$> index <*> pure stream
                      in take count <$> nthSet

mkTable :: CharSet -> CipherTable
mkTable charset = let generator = (`genRow` charset)
                      strings   = mapM generator charset
                  in createMap $ (zip charset) <$> strings
  where createMap Nothing = M.empty
        createMap (Just x) = M.fromList x

-- generates a cryptographic key based on a message and a keyword
-- if the message is "hello there", and the keyword is "apple";
-- then the resulting key will be "appleapplea"
genKey :: String -> String -> String
genKey msg kw = kw `padTo` length msg
  where padTo xs size = take size $ cycle xs

-- gets a row from the table at a specific key k
-- if no row exists, Nothing is returned
getRow ::  Char -> CipherEnv -> Maybe String
getRow k e = M.lookup k $ cipherTable e

-- gets the vth letter on the kth row on the table
-- due to the fact that both getRow and baseIndex are Maybe types,
-- Applicative mapping needs to be used
encryptLetter :: Crypt Char -- table[k, base[v]]
encryptLetter k v e = (!!) <$> getRow k e <*> baseIndex
  where baseIndex = L.elemIndex v base -- gets index from base
        base = cipherChars e

-- gets the index of the cth letter on the kth row in the table,
-- then uses that index to find the cth letter in the base
decryptLetter :: Crypt Char -- base[table[k, c]]
decryptLetter k c e = (base !!) <$> tableIndex
  where tableIndex = getRow k e >>= L.elemIndex c -- gets index from table
        base = cipherChars e

-- maps a function to two strings, making sure both strings are up-cased
-- mapM is used because f returns a monad,
-- and neither of the strings are monads themselves
crypt :: Crypt Char -> Crypt String
crypt f w kw e = mapM (uncurry3 f) zipped
  where word = upCase w
        keyword = upCase kw
        zipped = zip3 key word $ repeat e
        key = genKey word keyword

-- encrypts the string using a message and a keyword
encrypt :: Crypt String
encrypt = crypt encryptLetter

-- decrypts using a message and a keyword
decrypt :: Crypt String
decrypt = crypt decryptLetter

-- helper function to turn a string into upper-case
upCase :: String -> String
upCase = map toUpper

uncurry3 :: (a -> b -> c -> d) -> (a, b, c) -> d
uncurry3 f = \(a, b, c) -> f a b c

mkEnv :: CharSet -> CipherEnv
mkEnv c = Cipher c $ mkTable c

-- Runs a function with an environtment and inputs
-- example: runEnv ['a' .. 'z'] encrypt "helloworld" "cazoo"
runEnv :: CharSet -> Crypt String -> String -> String -> Maybe String
runEnv c f a b =  f a b $ mkEnv c

-- Similar to runEnv, except lets the user use the functions directly
-- example: runFun ['a' .. 'z'] $ encrypt "helloworld" "cazoo"
runFun :: CharSet -> (CipherEnv -> Maybe String) -> Maybe String
runFun c f = f $ mkEnv c

Here's an example main function that runs the code:

main :: IO ()
main = do
  let characterSet = ' ':',':['A' .. 'Z'] ++ ['a' .. 'z']
      message      = "Hello, this is a test"
      keyword      = "bananas"
      environment  = runFun characterSet
      cipher       = environment $ encrypt message keyword
  putStrLn $ "Here's the output of running \"" ++ message ++ "\" through the function: "
  putStrLn $ show cipher

Which should output:

Here's the output of running "Hello, this is a test" through the function: 
Just "KGaNdBSWJXUNKmBCNVTUn"
\$\endgroup\$
9
\$\begingroup\$

removing the limit of just having letters from A to Z, and nothing more.

You didn't. And that's due to upCase and some quirks with CharSet. But before we get to that and the actual code review, let's have a look how you could have tested that:

-- add to end of main
  putStrLn $ "Here's the output of running " ++ show cipher ++ " back through the function: "
  putStrLn $ show (cipher >>= \m -> environment $ decrypt m keyword)

Result:

Here's the output of running "Hello, this is a test" through the function:
Just "KGaNdBSWJXUNKmBCNVTUn"
Here's the output of running Just "KGaNdBSWJXUNKmBCNVTUn" back through the function:
Just "HEnLq, THIS Iu A TESv"

That doesn't seem right. The wrong case of E in "Hello" indicates that there's something going on with casing, and indeed, if you change upCase to id, your cipher will work. You could even add some QuickCheck tests:

quickCheck $ property $
   forAll (listOf1 arbitrary)          $ \charset ->
   forAll (listOf1 $ elements charset) $ \message ->
   forAll (listOf1 $ elements charset) $ \keyword ->
     let run f m   = runFun charset $ f m keyword
         encrypted = run encrypt message
         decrypted = encrypted >>= run decrypt
     in maybe "***ERROR***" id decrypted == message

Which will show you the second error in your suite (or at least wrong assumptions to the CharSet):

*** Failed! Falsifiable (after 20 tests):
"\NAK\236j\NAKnL"                            [remark: see the duplicate \NAK?]
"\NAKL\NAKjj\236L\NAK\NAKL\236jL\236\NAK"
"\NAKL\NAK\NAK\236Ln\NAKLn\NAKjj"

The tests will pass if we use unique symbols in the character set:

forAll (L.nub `fmap` listOf1 arbitrary) $ \charset ->

And now, we finally get to code review.

No explicit export lists

module Vigenere where

This will export all your methods to the client. While this gives them much power, it also fills their namespace with all your helper methods, like genKey, upCase and others. Instead, export only the functions they're going to need. If you want to give them access to internal methods, split your module into Vigenere and Vigenere.Internal.

That will keep the noise down for common users, but give full power to developers.

Assumptions on CharSet do not hold

Let's have a look at your types:

type Crypt a = a -> a -> CipherEnv -> Maybe a

You should definitely explain what a means here. Especially what the order of arguments is, e.g. "A Crypt takes a message, a key, and a cipher environment and produces an encrypted".

type CharSet = [Char] -- alias for string, used to distinguish from text input
type CipherTable = M.Map Char CharSet

The word "alias" together with "distinguish" should ring some bells. While it's nice to see CharSet in the type signature, it won't hinder a user to feed plain Strings into your cipher, which completely destroys an assumption you hold throughout your program:

A CharSet does not contain any symbol twice.

Now, you can either specify that CharSet must not contain duplicate symbols, or you use a newtype with a smart constructor:

newtype CharSet = CharSet String 

-- | Creates a CharSet from a String.
--   It basically removes all duplicate symbols.
mkCharSet :: String -> CharSet
mkCharSet = CharSet . L.nub

Note that this also contains a lot potential for future optimizations, e.g. using a sorted Data.Vector in order to have a O(log(N)) lookup instead of O(N). Note that you would not export the CharSet data constructor in this case, only mkCharSet (see previous section).

That being said, the real culprit is hidden in genRow. CharSet could be an alias to String without smart constructor, but your use in genRow doesn't allow that.

Use assumptions to make functions easier

Your genRow is a little bit odd. It looks for a character, drops all other till it finds it, and then adds them at the end. That's not the odd part. The odd part is it returns a Maybe, although all invariants in mkTable—the only place genRow is used—make sure that char will always be contained in charset.

We can simplify this function a lot if we use break instead:

-- | Generates the shifted 'CharSet' corresponding
--   to the given 'char'. Returns the 'CharSet'  if
--   'char' is not contained.
genRow :: Char -> CharSet -> CharSet
genRow char charset = 
   let (a, b) = break (== char) charset
   in b ++ a

This is a function you wouldn't likely export. Now, mkTable can also get simplified a lot:

-- | Generates a CipherTable from a 'CharSet'.
--   Note that decrypting will fail if the
--   'CharSet' contains duplicate symbols.
mkTable :: CharSet -> CipherTable
mkTable charset = M.fromList $ map generator charset
  where generator c = (c, genRow c charset)
  --                  ^^^^^^^^^^^^^^^^^^^^^

It should now be obvious that the aforementioned assumptions about genRow's arguments really hold.

Use Haddock documentation syntax

I suck at this, so I'm a little bit of a hypocrite here. However, all following functions have documentation (great!), so it would be even better if the documentation could get exported via cabal haddock (or similar commands).

If you do export documentation, make sure that implementation details aren't in there, e.g.

-- due to the fact that both getRow and baseIndex are Maybe types,
-- Applicative mapping needs to be used

is nice to know for you, but not really necessary for a user.

Make crypt a little bit simpler

crypt :: Crypt Char -> Crypt String
crypt f w kw e = sequence $ zipWith3 f key w (repeat e)
  where key = genKey w kw

Remember that upCase was causing trouble? This was the place where you've used it. zipWith3 enables you to skip the intermediate list of triples. That removes the need of uncurry3 in your code.

Note that this could get a little bit more optimized if you don't repeat the environment, but fuse it with the function:

crypt :: Crypt Char -> Crypt String
crypt f w kw e = sequence $ zipWith f' key w
  where key    = genKey w kw
        f' a b = f a b e

This would be a lot easier if Crypt was different.

Change the order of types in Crypt

There's one thing I dislike throughout your code, and I'm sorry for that. But Crypt isn't really easy to use. That's because the environment, the one thing that most likely won't change in consecutive calls, comes last in your code.

If I want to encrypt ASCII messages, I have to write

encryptASCII :: String -> Maybe String
encryptASCII m w = crypt encryptLetter m w (mkEnv ['\NUL'..'\DEL'])

But if you have a look at quickCheckWith or other similar *With functions, you'll notice that the additional parameters come first. If this was the case in crypt (and therefore Crypt), I could write

encryptASCII = crypt encryptLetter (mkEnv ['\NUL'..'\DEL'])
decryptASCII = crypt decryptLetter (mkEnv ['\NUL'..'\DEL'])

Your functions runEnv and runFun would also get a lot easier. However, as I said, this is mostly opinion based. But think about it: if your encrypting and decrypting, your usually staying in the same CipherEnv. And if you're encrypting multiple messages (for example for parallel encryption), you also use the same key.

For an inspiration, have a look at Socket. See that Socket is always first? This enables one to easily write things like

mapM_ (send socket) data

If send's type was send :: String -> Socket -> IO Int, one would have to use

mapM_ (flip send socket) data
-- or
mapM_ (\s -> send s socket) data

So arguments that change "less often" (e.g. stay the same over multiple applications) should be first.

Summary

Overall, good work. Beside the two errors, the Vigenère Cipher works as intended. So, what's left to do?

  • Restrict your exports
  • Add tests:
    • Check that for a single character key your Vigenère works the same as Caesar.
    • Check that encrypted text can get decrypted (if both return Just ...)
    • Check what happens if you supply symbols that are not in the CharSet
  • Add documentation
  • Think about Crypt, and whether you want to change the order of arguments.
\$\endgroup\$
  • \$\begingroup\$ I admit, the upcase was a leftover from the original version that didn't allow custon charsets. I forgot to remove it \$\endgroup\$ – Electric Coffee Dec 27 '15 at 13:24
  • \$\begingroup\$ could you add an example of how using it as a reader could be beneficial? Also I'm not sure what exactly you meant about Crypt's argument order being wacky; could you elaborate on that? I'm learning a lot from your examples so far, so thanks \$\endgroup\$ – Electric Coffee Dec 27 '15 at 13:47
  • \$\begingroup\$ @ElectricCoffee: Sorry, that will have to wait for some hours, at least for the inclusion in the answer. Here is some preliminary answer: the Reader part is more or less a piece of trivia. You can write monadic code instead of your usual one, if you want to, which is feasible sometimes. The "wacky argument order" is a opinion, but if you use (pseudo) Crypt = environment -> key -> message -> Maybe cipher, one can easily create own encryption/decryptions via encrypt = crypt encryptLetter someCharSet, decrypt = decryptLetter someCharset. Note that I was able to use point-free style here. \$\endgroup\$ – Zeta Dec 27 '15 at 18:35
  • \$\begingroup\$ I tried following your suggestion and implemented it using Reader. But due to a lack of ideas as to how to do it properly, every function ending in CipherEnv -> Maybe a has instead become ReaderT CipherEnv Maybe a, which indeed for better or worse let me write all the code in a monadic style. I've pastebinned the code here if you're interested. \$\endgroup\$ – Electric Coffee Dec 27 '15 at 20:05
  • \$\begingroup\$ @ElectricCoffee: I removed the Reader remark, since it was slightly misleading. Sorry for the confusion. \$\endgroup\$ – Zeta Dec 28 '15 at 9:19
9
\$\begingroup\$

Thanks to Zeta work, my answer will be shorter :-)

If you look at Wikipedia, especially the Algebraic description of the Vigenère cipher, there is a quicker way to encrypt/decrypt.

Using a bidirectional map

This method requires the use of a bidirectional map. Such a package is available on Hackage (Data.Bimap) and it only requires the containers package.

In pseudo language:

encrypted[i] = (input[i] + password[i]) mod alphabetSize
decrypted[i] = (input[i] - password[i]) mod alphabetSize

Note:

  • these computations are made on integers,
  • i is the index of a character from 0 to length of the input,
  • password is considered cycled.

This said, we need a structure which can convert from a to Int and back from Int to a: the Bimap type.

import Data.Bimap

type Alphabet a = Bimap Int a

mkAlphabet :: Ord a => [a] -> Alphabet a
mkAlphabet = foldl (flip =<< insert . size) empty

The mkAlphabet, given a list of elements, will generate a bidirectional map between Int and these elements:

alphabet = mkAlphabet ['A'..'Z']

alphabet ! 3 -- 'D'
alphabet !> 'D' -- 3
size alphabet -- 26

In the Vigenère cipher, characters are independent from each other. We can therefore write a function working on one character and its corresponding character from the password. This is a simple transcription from the pseudo language functions presented before:

-- a: Alphabet
-- c: character to encrypt
-- k: corresponding character from the password
vigenereEnc a c k = a ! ((a !> c + a !> k) `mod` size a)

-- a: Alphabet
-- d: character to decrypt
-- k: corresponding character from the password
vigenereDec a d k = a ! ((a !> d - a !> k) `mod` size a)

We can now apply these functions to a list of characters. The zipWith function is handy for this:

encrypt :: Ord a => Alphabet a -> [a] -> [a] -> [a]
encrypt alphabet s key = zipWith (vigenereEnc alphabet) s (cycle key)

decrypt :: Ord a => Alphabet a -> [a] -> [a] -> [a]
decrypt alphabet s key = zipWith (vigenereDec alphabet) s (cycle key)

Assembling everything

Our Vigenère module looks like this after assembled:

module Vigenere (mkAlphabet, encrypt, decrypt) where

import Data.Bimap (size, (!), (!>), Bimap, insert, size, empty)

type Alphabet a = Bimap Int a

mkAlphabet :: Ord a => [a] -> Alphabet a
mkAlphabet = foldl (flip =<< insert . size) empty

encrypt :: Ord a => Alphabet a -> [a] -> [a] -> [a]
encrypt alphabet s key = zipWith (vigenereEnc alphabet) s (cycle key)
    where vigenereEnc a c k = a ! ((a !> c + a !> k) `mod` size a)

decrypt :: Ord a => Alphabet a -> [a] -> [a] -> [a]
decrypt alphabet s key = zipWith (vigenereDec alphabet) s (cycle key)
    where vigenereDec a d k = a ! ((a !> d - a !> k) `mod` size a)

And the Main module:

module Main where

import Vigenere (mkAlphabet, encrypt, decrypt)

main :: IO ()
main = do
    let alphabet = mkAlphabet $ ' ':',':['A' .. 'Z'] ++ ['a' .. 'z']
        message  = "Hello, this is a test"
        password = "bananas"

        cipher   = encrypt alphabet message password
        decipher = decrypt alphabet cipher password

    putStrLn $ "Clear: " ++ message
    putStrLn $ "Encrypted: " ++ cipher
    putStrLn $ "Decrypted: " ++ decipher

Notes

This version does not handle the case where the character being read is not in the alphabet. It’s not a wrong way to do things: if you look at the division operator, it’s up to the caller to take care of a division by zero, leaving it free to handle this special case whatever way it thinks is good.

In our case, the function will not need to be rewritten with each of the scenarios:

  • stops at the first unknown character and output an error,
  • stops at the first unknown character and output the encrypted string so far
  • ignore all unknown characters.

This function is also generic: you can work with whatever type you want as long as it supports the Ord class.

\$\endgroup\$
  • \$\begingroup\$ While this is an (valid) complete other way to achieve a Vigenère en-/decryption, doesn't your post, well, … essentially propose to rewrite it? I'm new to SE.CR, so this might be feasible, but I thought that the point of CR would be to save as much as possible from the original code, and only propose a complete other approach as a last resort. That being said, on StackOverflow, every answer should stand for itself. You probably want to mention things like limited exports and so on (and a newtype instead of type, maybe). \$\endgroup\$ – Zeta Dec 27 '15 at 18:28
  • \$\begingroup\$ Sorry if I gave you the impression that everything should be rewritten. My answer focused on the "I feel like I way overdid it" part of the post. I was aiming to show that using adequate types greatly simplify the code. The code is here only to show how this can be achieved (keeping in sight the variable cipher table), not to say "the vigenere module must be written like this". \$\endgroup\$ – zigazou Dec 27 '15 at 21:07

Your Answer

By clicking “Post Your Answer”, you agree to our terms of service, privacy policy and cookie policy

Not the answer you're looking for? Browse other questions tagged or ask your own question.