Solution of the hacker rank expressions problem in Haskell

HackerRankExpressions.hs

{-# LANGUAGE RecordWildCards #-}

-- | Problem https://www.hackerrank.com/challenges/expressions

import Control.Applicative
import Control.Arrow
import Control.Monad
import Control.Monad.Trans.Class
import Control.Monad.Trans.Maybe
import Control.Monad.Trans.State
import Data.List
import Data.Maybe
import Data.Set (Set)
import qualified Data.Set as Set

-- | Operator data type and related functions.
data Op = Add | Sub | Mul
  deriving (Show, Eq)

allOps :: [Op]
allOps = [Add, Sub, Mul]

showOp :: Op -> String
showOp Add = "+"
showOp Sub = "-"
showOp Mul = "*"

evalOp :: Op -> Int -> Int -> Int
evalOp Add = (+)
evalOp Sub = (-)
evalOp Mul = (*)

-- | Expression data type and related functions.
data Expr = Val Int | App Expr Op Int -- left-associative application
  deriving (Show)

showExpr :: Expr -> String
showExpr (Val x) = show x
showExpr (App e op x) = showExpr e ++ showOp op ++ show x

evalModExpr :: Int -> Expr -> Int
evalModExpr m (Val x) = x `mod` m
evalModExpr m (App e op x) = evalOp op (evalModExpr m e) x `mod` m

-- | Tree data type to capture all possible expressions, including partial ones.
-- The level field is used to keep track of the depth of the tree.
type Level = Int

data Tree = Tree
  { level :: Level,
    expr :: Expr,
    edges :: [Tree]
  }
  deriving (Show)

-- | Build a tree of all possible expressions given a list of integers.
buildTree :: [Int] -> Tree
buildTree (x : xs) =
  Tree 0 (Val x) [go 1 (Val x) op xs | op <- allOps]
  where
    go :: Level -> Expr -> Op -> [Int] -> Tree
    go l e o [x] = Tree l (App e o x) [] -- leaf node
    go l e o (x : xs) = Tree l (App e o x) [go (l + 1) (App e o x) o' xs | o' <- allOps]

-- | Find the expression that evaluates to a multiple of 101.
-- | Cache is used to avoid duplicate calculations.
-- | On a a given level, if the same value is encountered, we can skip the calculation.
type Cache = Set (Level, Int)

-- | Monad stack to handle state and maybe computations.
-- | MaybeT is used to short-circuit the computation when a solution is found.
-- | State needs to be the inner monad so that the cache can be updated even when the computation fails.
type MonadStack a = MaybeT (State Cache) a

runMonadStack :: MonadStack a -> a
runMonadStack = fromJust . flip evalState Set.empty . runMaybeT

-- | Similar to foldMap, but for Alternative instances.
-- | This is used to short-circuit the computation when a solution is found.
foldMapA :: (Foldable t, Alternative f) => (a -> f b) -> t a -> f b
foldMapA f = foldr ((<|>) . f) empty

-- | Helper function to find the expression that evaluates to a multiple of 101.
-- | The function uses MonadStack to handle state and maybe computations.
findExprHelper :: Tree -> MonadStack Expr
findExprHelper Tree {..} = do
  let v = evalModExpr 101 expr
  let k = (level, v)
  s <- lift get
  -- Skip if the same value is encountered.
  guard $ not $ k `Set.member` s
  -- Update the cache.
  lift . modify $ Set.insert k
  if v == 0 && null edges
    -- Found a solution.
    then pure expr
    -- Continue searching next level.
    else foldMapA findExprHelper edges

-- | Find the expression that evaluates to a multiple of 101.
findExpr :: Tree -> Expr
findExpr = runMonadStack . findExprHelper

-- | Top level functions to parse, solve, and print the solution.
parse :: String -> [Int]
parse = lines >>> last >>> words >>> map read

solve :: [Int] -> Expr
solve = buildTree >>> findExpr

main :: IO ()
main = interact $ showExpr . solve . parse