Parsing a simple AST with Prisms (without detailed syntax error reporting)

ast-parsing-with-prisms.hs

{-# LANGUAGE RankNTypes, TupleSections, TemplateHaskell, DeriveGeneric #-}

import Control.Applicative
import Control.Lens
import qualified Data.Char as Char
import GHC.Generics
import Test.QuickCheck
import Test.QuickCheck.Arbitrary.ADT

data Expr
    = Lit Int
    | Add Expr Expr
    | Mul Expr Expr
    deriving (Eq, Show, Generic)

makePrisms ''Expr

expr :: Prism' String Expr
expr =
    tokens .      -- convert string to tokens
    takeExpr .    -- parse an expression
    secondOnly [] -- there should be no remaining tokens

takeExpr :: Prism' [String] (Expr, [String])
takeExpr =
    infixOpLeftRecursion "+" _Add $ -- Additions of
    infixOpLeftRecursion "*" _Mul $ -- multiplications of
    tryMatch (asideFirst _Lit) (_Cons . asideFirst _Show) $ -- Literals or
    _Cons . firstOnly "(" . takeExpr . aside (_Cons . firstOnly ")") -- stuff in parenthesis

infixOpLeftRecursion ::
    Eq a =>
    a ->
    Prism' expr (expr, expr) ->
    Prism' [a] (expr, [a]) ->
    Prism' [a] (expr, [a])
infixOpLeftRecursion operatorText cons sub =
    leftRecursion cons
    (aside (_Cons . firstOnly operatorText . sub) . retuple)
    sub

leftRecursion ::
    Prism' whole cons ->
    Prism' (whole, state) (cons, state) ->
    Prism' state (whole, state) ->
    Prism' state (whole, state)
leftRecursion cons extend base =
    prism' build (fmap parseExtends . (^? base))
    where
        build (x, state) =
            maybe (base # (x, state)) (build . (extend #) . (, state)) (x ^? cons)
        parseExtends x =
            x ^? extend <&> _1 %~ (cons #) & maybe x parseExtends

tryMatch ::
    Prism' whole cons ->
    Prism' src cons ->
    Prism' src whole ->
    Prism' src whole
tryMatch cons parse fallback =
    prism' build (\x -> (x ^? parse <&> (cons #)) <|> x ^? fallback)
    where
        build x = maybe (fallback # x) (parse #) (x ^? cons)

tokens :: Iso' String [String]
tokens =
    iso splitTokens (foldr addToken "")
    where
        addToken x "" = x
        addToken [x] y
            | Char.generalCategory x == Char.OpenPunctuation = x : y
        addToken x (y:ys)
            | Char.generalCategory y == Char.ClosePunctuation = x <> (y:ys)
        addToken x y = x <> " " <> y
        isOp = (`elem` [Char.MathSymbol, Char.OtherPunctuation]) . Char.generalCategory
        isParen = (`elem` "()[]{}")
        splitTokens "" = []
        splitTokens (x:s:xs) | Char.isSpace s = [x] : splitTokens xs
        splitTokens (s:xs) | Char.isSpace s = splitTokens xs
        splitTokens (x:xs) | isParen x = [x] : splitTokens xs
        splitTokens (x:xs) =
            case splitTokens xs of
            [] -> [[x]]
            ((y:ys) : zs) | not (isParen y) && isOp x == isOp y -> (x:y:ys) : zs
            ys -> [x] : ys

asideFirst :: APrism s t a b -> Prism (s, e) (t, e) (a, e) (b, e)
asideFirst l = swapped . aside l . swapped

firstOnly :: Eq e => e -> Prism' (e, a) a
firstOnly x = asideFirst (only x) . iso snd ((,) ())

secondOnly :: Eq e => e -> Prism' (a, e) a
secondOnly x = swapped . firstOnly x

retuple ::
    Iso
    (a0, (a1, a2)) (b0, (b1, b2))
    ((a0, a1), a2) ((b0, b1), b2)
retuple =
    iso
    (\(w0, (w1, r)) -> ((w0, w1), r))
    (\((w0, w1), r) -> (w0, (w1, r)))

instance Arbitrary Expr where
    arbitrary = genericArbitrary
    shrink = genericShrink

propParseBack :: Expr -> Bool
propParseBack e = (expr # e) ^? expr == Just e