CLI helper for discrete math, but in haskell

discrete.hs

import Prelude hiding (gcd, lcm)
import Data.Bits ((.&.), Bits)
import Data.Semigroup
import System.Exit (die)
import System.Environment (getProgName, getArgs)
import Text.Printf (printf)

-- | Calculate whether a number is odd
isOdd :: (Integral a, Data.Bits.Bits a) => a -> Bool
isOdd n = n .&. 1 == 1

-- | Given a list of values, return a list of tuples (index, value)
enumerate :: [a] -> [(Integer, a)]
enumerate = zip [0..]

factorial :: Integral a => a -> a
factorial n = product [1..n]

-- | Calculate C(n, r), n choose r
ncr :: Integral a => a -> a -> a
ncr n r = if r > n
  then 0
  else let r' = min r (n-r) in -- small optimization
    -- n! / (r! * (n-r)!) == n*(n-1)*..(n-r+1) / r!
    let numerator = foldl (*) 1 [n,n-1..n-r'+1]
        denominator = factorial r'
    in numerator `div` denominator

-- | Calculate a sterling number of the second kind:
-- | the number of ways to partition an n-element set into k non-empty sets
sterling :: Integral a => a -> a -> a
sterling n 1 = 1
sterling n k = if n < 1 || k > n then 0
  else if n == k then 1
  else sterling (n-1) (k-1) + k*sterling (n-1) k

-- | Calculate the ways to partition a number n into k positive terms
partition :: Integral a => a -> a -> a
partition n k = if k <= 0 || k > n then 0
  else if k == n then 1
  else partition (n-k) k + partition (n-1) (k-1)

-- | Calculate the greatest common divisor of n and d using the Euclidean algorithm
gcd :: Integral a => a -> a -> a
gcd n 0 = n
gcd n d = gcd d (n `mod` d)

-- | Find the least common multiple of n and arbitrarily many other numbers
lcm :: Integral a => [a] -> a
lcm [] = 0
lcm [n] = n
lcm [n, m] = n*m `div` gcd n m
lcm (n:m:ns) = lcm (lcm [n, m]:ns)

-- | Given an element k in the group Z mod n, find the order of k
order :: Integral a => a -> a -> a
order k n = lcm [k, n] `div` k

-- | Calculate whether two numbers are relatively prime,
-- | i.e. have no common factors other than 1
coprime :: Integral a => a -> a -> Bool
coprime n k = gcd n k == 1

-- | Given the order of an arbitrary number of groups G_i,
-- | find whether their direct product G is cyclic
-- Theorem: G is cyclic <=> order(G_i) and order(G_j) are relatively prime
-- for all i != j
-- Theorem: a and b are relatively prime <=> gcd(a, b) == 1
cyclic :: Integral a => [a] -> Bool
cyclic orders = all (\(i, order_i) ->
  all (\(j, order_j) ->
    i == j || coprime order_i order_j
  ) $ enumerate orders) (enumerate orders)

-- | Given a set S of n elements, calculate the number of permutations of S
-- | such that no element remains in its original position
-- | NOTE: this is slower but more accurate than n! / e
-- Algorithm: D(n) = \sum_{i=0}^{n} (-1)^i n! / i!
--                 = n! * \sum_{i=0}^{n} (-1)^i 1 / i!
-- NOTE: it is a logic error for this to be anything other than an Integer
derangementS :: Integer -> Integer
derangementS n = id
  . sum
  . zipWith (*) signs
  . scanl (*) 1
  $ [n,n-1..1]
  where signs = drop (fromIntegral (n .&. 1)) (cycle [1,-1])

-- see https://en.wikipedia.org/wiki/Derangement#Counting_derangements
-- where it says !n = n[!(n-1)] + (-1)^n
derangements :: Integral a => [a]
derangements = scanl (\a (b,c) -> a*b+c) 1 . zip [1..] $ cycle [-1,1]

{- the rest of the file is _very_ non-idiomatic for haskell,
 - because I wanted to do something it's not designed for.
 -
 - The original python code looked something like this:
 - { "myfunc1": myfunc, "myfunc2": myfunc }[argv[0]](*argv[1:])
 -
 - Haskell clearly does not allow functions of different types
 - to be stored in the same data structure nor for functions to be called
 - with a variable number of arguments.
 - Instead, I wrote a wrapper around the type system to allow arbitrary arguments
 - and arbitrary function outputs so all the functions would have the same type.
 -
 - Many thanks to @jle` and @lyxia for the help on IRC!
 -}
data FuncOutput = I Integer | B Bool
data FuncResult = Ok FuncOutput | NameErr | ArityErr Integer

resultOf :: String -> [Integer] -> FuncResult
resultOf = applyFuncs
  [ Func ["D", "derangement", "derangements"] (Arity1 (\n -> I $ derangements !! fromInteger n))
  , Func ["!", "factorial"] (Arity1 (I . factorial))
  , Func ["g", "gcd", "euclidean"] (Arity2 (\x y -> I (gcd x y)))
  , Func ["p", "P", "partition"] (Arity2 (\x y -> I (partition x y)))
  , Func ["s", "S", "sterling"] (Arity2 (\x y -> I (sterling x y)))
  , Func ["order"] (Arity2 (\x y -> I $ order x y))
  , Func ["c", "C", "ncr", "combinations"] (Arity2 (\x y -> I (ncr x y)))
  , Func ["coprime", "relativelyprime"] (Arity2 (\x y -> B (coprime x y)))
  , Func ["lcm"] (ArityArbitrary (I . lcm))
  , Func ["cyclic"] (ArityArbitrary (B . cyclic))
  ]

data Func = Func
  [String]  -- names
  FuncBody  -- function definition

data FuncBody
  = Arity1 (Integer -> FuncOutput)
  | Arity2 (Integer -> Integer -> FuncOutput)
  | ArityArbitrary ([Integer] -> FuncOutput)

arity :: FuncBody -> Integer
arity (Arity1 _) = 1
arity (Arity2 _) = 2

applyFunc :: Func -> String -> [Integer] -> FuncResult
applyFunc (Func names body) n args
  | n `elem` names = case (body, args) of
      (Arity1 f, [n]) -> Ok (f n)
      (Arity2 f, [n, k]) -> Ok (f n k)
      (ArityArbitrary f, _) -> Ok (f args)
      (body, _) -> ArityErr (arity body)
  | otherwise = NameErr

-- for foldMap
instance Semigroup FuncResult where
  Ok i <> _ = Ok i
  ArityErr i <> _ = ArityErr i
  NameErr <> j = j

instance Monoid FuncResult where
  mempty = NameErr
  mappend = (<>)

applyFuncs :: [Func] -> String -> [Integer] -> FuncResult
applyFuncs = foldMap applyFunc

variables = ["n", "m", "k", "l", "p", "q"]

main = do
  self <- getProgName
  args <- getArgs
  let nums = map read args
  case resultOf self nums of
    Ok (I x) -> print x
    Ok (B x) -> print x
    ArityErr i -> printf "usage: %s %s\n" self $ unwords (take (fromInteger i) variables)
    NameErr -> putStrLn $ "unknown function " ++ self