megaparsec with custom token stream (Lib.hs -- general definitions, Simple.hs & Lam.hs -- examples of usage)

Lam.hs

{-# LANGUAGE FlexibleInstances #-}
{-# LANGUAGE LambdaCase #-}
{-# LANGUAGE OverloadedStrings #-}

module Lam where

import Control.Applicative
import qualified Control.Monad.Combinators.Expr as Expr
import qualified Data.Char as Char
import qualified Data.List.NonEmpty as NE
import Data.Text (Text)
import qualified Data.Text as T
import Data.Void (Void)
import Text.Megaparsec (Parsec)
import qualified Text.Megaparsec as MP
import qualified Text.Megaparsec.Char as MPC
import qualified Text.Megaparsec.Char.Lexer as L

import Lib

testParse :: (Show a) => Parser a -> Text -> IO ()
testParse p txt = case MP.runParser lexer "" txt of
  Left e -> putStrLn $ "Lexing error:\n" <> MP.errorBundlePretty e
  Right ts -> case MP.runParser p "" (LexicalStream txt ts) of
    Left e -> putStrLn $ "Parsing error:\n" <> MP.errorBundlePretty e
    Right a -> print a

main :: IO ()
main = do
  let test txt = do
        putStrLn $ T.unpack txt
        testParse (pExp <* MP.hidden MP.eof) txt
  test ""
  test "a + b * c"
  test "\\f x y -> f y x"
  test "a + -b"
  test "\"xyz\" + uv"
  test "foo // bar"

--- AST

data Exp
  = ELit Lit
  | EVar Ident
  | EApp Exp Exp
  | ELam [Ident] Exp
  | EUn Op Exp
  | EBin Exp Op Exp
  deriving (Show, Eq, Ord)

data Lit
  = LInt (WithOffset Int)
  | LStr (WithOffset Text)
  deriving (Show, Eq, Ord)

type Ident = WithOffset Text

type Op = WithOffset Text

--- Tokens

data Tok
  = TLitInt Int
  | TLitStr Text
  | TIdent Text
  | TBinOp Text
  | TUnOp Text
  | TLPar -- (
  | TRPar -- )
  | TLam -- \
  | TArr -- ->
  deriving (Show, Eq, Ord)

prettyTok :: Tok -> String
prettyTok = \case
  TLitInt n -> "integer " <> show n
  TLitStr s -> "literal " <> "\"" <> T.unpack s <> "\""
  TIdent n -> "identifier " <> T.unpack n
  TUnOp o -> "unary operation " <> T.unpack o
  TBinOp o -> "binary operation " <> T.unpack o
  TLPar -> "opening parenthesis"
  TRPar -> "closing parenthesis"
  TLam -> "lambda"
  TArr -> "arrow"

tokenLength :: Tok -> Int
tokenLength = \case
  TLitInt n -> length $ show n
  TLitStr s -> 2 + T.length s
  TIdent n -> T.length n
  TUnOp o -> T.length o
  TBinOp o -> T.length o
  TLPar -> 1
  TRPar -> 1
  TLam -> 1
  TArr -> 2

instance MP.VisualStream [WithOffset Tok] where
  showTokens _ = unwords . NE.toList . fmap (prettyTok . unOffset)
  tokensLength _ ts = o1 - o0 + tokenLength t1
    where
      WithOffset o0 _ = NE.head ts
      WithOffset o1 t1 = NE.last ts

-- Lexer

type Lexer = Parsec Void Text

sc :: Lexer ()
sc = MPC.space

lexeme :: Lexer a -> Lexer (WithOffset a)
lexeme = L.lexeme sc . withOffset

symbol :: Text -> Lexer (WithOffset Text)
symbol = lexeme . MPC.string

lexer :: Lexer [WithOffset Tok]
lexer = MP.manyTill lTok $ MP.hidden MP.eof

lTok :: Lexer (WithOffset Tok)
lTok =
  MP.choice
    [ (TLPar <$) <$> symbol "(",
      (TRPar <$) <$> symbol ")",
      (TLam <$) <$> symbol "\\",
      (TArr <$) <$> symbol "->",
      fmap TLitInt <$> lLitInt,
      fmap TLitStr <$> lLitStr,
      fmap TIdent <$> lIdent,
      fmap TUnOp <$> lUnOp,
      fmap TBinOp <$> lBinOp
    ]

lLitInt :: Lexer (WithOffset Int)
lLitInt = lexeme L.decimal

lLitStr :: Lexer (WithOffset Text)
lLitStr = lexeme $ fmap T.pack lRawStr
  where
    lRawStr = MPC.char '\"' *> MP.manyTill L.charLiteral (MPC.char '\"')

lIdent :: Lexer (WithOffset Text)
lIdent =
  lexeme $
    T.cons <$> MP.satisfy headChar <*> MP.takeWhileP Nothing tailChar
  where
    headChar = Char.isAlpha
    tailChar c = Char.isAlphaNum c || c == '_'

lUnOp :: Lexer (WithOffset Text)
lUnOp = lexeme $ MP.try $ lOp_ <* MP.notFollowedBy MPC.spaceChar

lBinOp :: Lexer (WithOffset Text)
lBinOp = lexeme lOp_

lOp_ :: Lexer Text
lOp_ = MP.takeWhile1P Nothing opChar
  where
    opChar c = not $ Char.isSpace c || (c `elem` ['(', ')', '_']) || Char.isAlphaNum c

--- Parser

type Parser = Parsec Void (LexicalStream Text [WithOffset Tok])

pExp :: Parser Exp
pExp = Expr.makeExprParser pTerm ops
  where
    ops :: [[Expr.Operator Parser Exp]]
    ops =
      [ map un ["-", "!"],
        [Expr.InfixL $ pure EApp]
      ]
        <> map (map bin) [["*", "/"], ["+", "-"]]
    un s = Expr.Prefix $ EUn . (s <$) <$> offsetted (TUnOp s)
    bin s = Expr.InfixL $ flip EBin . (s <$) <$> offsetted (TBinOp s)

pTerm :: Parser Exp
pTerm =
  MP.choice
    [ offsetted TLPar *> pExp <* offsetted TRPar,
      ELit <$> pLit,
      EVar <$> pIdent,
      ELam <$> (offsetted TLam *> some pIdent <* offsetted TArr) <*> pExp
    ]

pLit :: Parser Lit
pLit = LInt <$> pIntLit <|> LStr <$> pStrLit

pTokMay :: (Tok -> Maybe a) -> Parser (WithOffset a)
pTokMay f = MP.token (traverse f) mempty

-- One could use prisms instead of writing `f`s manually:

pIntLit :: Parser (WithOffset Int)
pIntLit = pTokMay f MP.<?> "integer"
  where
    f (TLitInt n) = Just n
    f _ = Nothing

pStrLit :: Parser (WithOffset Text)
pStrLit = pTokMay f MP.<?> "string literal"
  where
    f (TLitStr s) = Just s
    f _ = Nothing

pIdent :: Parser Ident
pIdent = pTokMay f MP.<?> "identifier"
  where
    f (TIdent n) = Just n
    f _ = Nothing

pUnOp :: Parser Op
pUnOp = pTokMay f MP.<?> "unary operation"
  where
    f (TUnOp op) = Just op
    f _ = Nothing

pBinOp :: Parser Op
pBinOp = pTokMay f MP.<?> "binary operation"
  where
    f (TBinOp op) = Just op
    f _ = Nothing

Lib.hs

{-# LANGUAGE DeriveTraversable #-}
{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE PartialTypeSignatures #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE TypeApplications #-}
{-# LANGUAGE TypeFamilies #-}
{-# LANGUAGE TypeOperators #-}
{-# LANGUAGE UndecidableInstances #-}

module Lib where

import Data.Data
import qualified Data.List.NonEmpty as NE
import qualified Data.Set as Set
import Text.Megaparsec

data LexicalStream rs ts = LexicalStream
  { rawStream :: rs,
    tokStream :: ts
  }

instance forall rs ts. (Stream ts) => Stream (LexicalStream rs ts) where
  type Token (LexicalStream rs ts) = Token ts
  type Tokens (LexicalStream rs ts) = Tokens ts

  tokenToChunk Proxy = tokenToChunk @ts Proxy
  tokensToChunk Proxy = tokensToChunk @ts Proxy
  chunkToTokens Proxy = chunkToTokens @ts Proxy
  chunkLength Proxy = chunkLength @ts Proxy
  chunkEmpty Proxy = chunkEmpty @ts Proxy
  take1_ (LexicalStream rs ts) = fmap (LexicalStream rs) <$> take1_ ts
  takeN_ n (LexicalStream rs ts) = fmap (LexicalStream rs) <$> takeN_ n ts
  takeWhile_ f (LexicalStream rs ts) = LexicalStream rs <$> takeWhile_ f ts

instance forall rs ts. (VisualStream ts) => VisualStream (LexicalStream rs ts) where
  showTokens Proxy = showTokens @ts Proxy
  tokensLength Proxy = tokensLength @ts Proxy

class HasRawOffset t where
  getRawOffset :: t -> Int

instance
  forall rs ts.
  (HasRawOffset (Token ts), Stream ts, TraversableStream rs) =>
  TraversableStream (LexicalStream rs ts)
  where
  reachOffsetNoLine o' (PosState (LexicalStream rs ts) o sp tw pref) =
    PosState (LexicalStream rs' ts') o' sp' tw' pref'
    where
      ts' = dropN_ (o' - o) ts
      PosState rs' _ sp' tw' pref' = case take1_ ts of
        Nothing -> PosState rs 0 sp tw pref
        Just (t, _) ->
          let ro = getRawOffset t
              ro' = maybe (ro + length_ rs - 1) (getRawOffset . fst) $ take1_ ts'
           in reachOffsetNoLine ro' $ PosState rs ro sp tw pref

  reachOffset o' (PosState (LexicalStream rs ts) o sp tw pref) =
    (mstr, PosState (LexicalStream rs' ts') o' sp' tw' pref')
    where
      ts' = dropN_ (o' - o) ts
      (mstr, PosState rs' _ sp' tw' pref') = case take1_ ts of
        Nothing -> (Nothing, PosState rs 0 sp tw pref)
        Just (t, _) ->
          let ro = getRawOffset t
              ro' = maybe (ro + length_ rs - 1) (getRawOffset . fst) $ take1_ ts'
           in reachOffset ro' $ PosState rs ro sp tw pref

dropN_ :: (Stream s) => Int -> s -> s
dropN_ n s = case takeN_ n s of
  Just (_, s') -> s'
  Nothing -> snd $ takeWhile_ (const True) s

length_ :: forall s. (Stream s) => s -> Int
length_ s = chunkLength @s Proxy $ fst $ takeWhile_ (const True) s

data WithOffset t = WithOffset Int t
  deriving (Eq, Ord, Show, Functor, Foldable, Traversable)

unOffset :: WithOffset t -> t
unOffset (WithOffset _ t) = t

instance HasRawOffset (WithOffset t) where
  getRawOffset (WithOffset o _) = o

liftToken :: t -> WithOffset t
liftToken = WithOffset 0

withOffset :: (MonadParsec e s m) => m a -> m (WithOffset a)
withOffset la = getOffset >>= \o -> WithOffset o <$> la

offsetted :: (Token s ~ WithOffset t, MonadParsec e s m, Ord t) => t -> m (WithOffset t)
offsetted c = token test (Set.singleton . Tokens . NE.singleton . liftToken $ c)
  where
    test ot =
      if unOffset ot == c
        then Just ot
        else Nothing

Simple.hs

{-# LANGUAGE FlexibleInstances #-}

module Simple where

import Data.Data
import qualified Data.List.NonEmpty as NE
import qualified Data.Set as Set
import Data.Void (Void)
import Text.Megaparsec

import Lib

data MyToken
  = Int Int
  | Plus
  | Mul
  | Div
  | OpenParen
  | CloseParen
  deriving (Eq, Ord, Show)

instance VisualStream [WithOffset MyToken] where
  showTokens Proxy =
    unwords
      . NE.toList
      . fmap (showMyToken . unOffset)
  tokensLength _ ts = o1 - o0 + tokenLength t1
    where
      WithOffset o0 _ = NE.head ts
      WithOffset o1 t1 = NE.last ts

showMyToken :: MyToken -> String
showMyToken t = case t of
  (Int n) -> show n
  Plus -> "+"
  Mul -> "*"
  Div -> "/"
  OpenParen -> "("
  CloseParen -> ")"

tokenLength :: MyToken -> Int
tokenLength t = length $ showMyToken t

type MyStream = LexicalStream String [WithOffset MyToken]

type Parser = Parsec Void MyStream

pInt :: Parser Int
pInt = token test Set.empty <?> "integer"
  where
    test (WithOffset _ (Int n)) = Just n
    test _ = Nothing

pSum :: Parser (Int, Int)
pSum = do
  a <- pInt
  _ <- offsetted Plus
  b <- pInt
  eof
  return (a, b)

exampleStream :: String -> MyToken -> MyStream
exampleStream symb constr =
  LexicalStream
    { rawStream = "5 " <> symb <> " 623",
      tokStream =
        [ WithOffset 0 (Int 5),
          WithOffset 2 constr,
          WithOffset 4 (Int 623)
        ]
    }

main :: IO ()
main = do
  putStrLn "Plus:"
  parseTest pSum (exampleStream "+" Plus)
  putStrLn "Div:"
  parseTest pSum (exampleStream "/" Div)

Compare with https://markkarpov.com/tutorial/megaparsec.html#working-with-custom-input-streams