Internal.hs

{-# LANGUAGE UndecidableInstances #-}

-- | Internal module, no stability guarantees
module Yasi.Internal
  ( Segment (..),
    parseSegments,
    ipExpr,
    interpolator,
    Displayish (..),
    Stringish (..),
  )
where

import Control.Monad ((>=>))
import qualified Data.Char as C
import qualified Data.Text as T
import qualified Data.Text.Display as TD
import qualified Data.Text.Lazy as TL
import qualified Data.Text.Lazy.Builder as TLB
import GHC.Generics (Generic)
import qualified Language.Haskell.Meta.Parse as GhcHsMeta
import qualified Language.Haskell.TH.Lib as TH
import qualified Language.Haskell.TH.Quote as TH
import qualified Language.Haskell.TH.Syntax as TH

data Segment
  = Lit String
  | Exp String
  | Abs -- idea due to interpolate
  deriving (Show, Eq, Generic)

parseSegments :: MonadFail m => Char -> String -> m [Segment]
parseSegments c = fmap (group []) . go
  where
    -- ugly, but simple enough™
    go s
      | let (lit, rest) = span (/= c) s, not (null lit) = (Lit lit :) <$> go rest
      | s == "" = pure []
      | s == [c] = fail $ "should not end with single " <> [c]
      | _ : c' : rest <- s, c == c' = (Lit [c] :) <$> go rest
      | _ : '{' : rest <- s = case span (/= '}') rest of -- TODO smarter?
          (exp, '}' : rest) ->
            let seg = if exp == "" then Abs else Exp exp
             in (seg :) <$> go rest
          _ -> fail "missing closing bracket"
      | _ : v : rest' <- s,
        isVarStartChar v =
          let (vs, rest) = span isVarChar rest'
           in (Exp (v : vs) :) <$> go rest
      | otherwise = fail $ "invalid char after " <> [c]
    isVarStartChar v = C.isAscii v && C.isAlpha v
    isVarChar v = C.isAscii v && (C.isAlphaNum v || v == '_' || v == '\'')

    group ls [] = lit ls
    group ls (Lit l : ss) = group (l : ls) ss
    group ls (s : ss) = lit ls <> (s : group [] ss)
    lit [] = []
    lit ls = [Lit $ mconcat $ reverse ls]

ipExpr :: (TH.Exp -> TH.Exp) -> [Segment] -> TH.Q TH.Exp
ipExpr transform segs = do
  (ls, lams) <- go segs
  pure
    . lams
    . transform
    . TH.AppE (TH.VarE 'stringish)
    . foldr (flip TH.UInfixE (TH.VarE '(<>))) (TH.VarE 'mempty)
    $ ls
  where
    go = \case
      [] -> pure ([], id)
      s : ss -> prep (go ss) case s of
        Lit l -> (,id) <$> [|$(TH.stringE l) :: String|]
        Exp e -> do
          exts <- TH.extsEnabled
          exp <- case GhcHsMeta.parseExpWithExts exts e of
            Right e -> pure e
            Left (line, col, msg) ->
              fail . unlines $
                [ "Parse error at splice `" <> e <> "`:",
                  show line <> ":" <> show col <> ": " <> msg
                ]
          pure (exp, id)
        Abs -> do
          n <- TH.newName "int"
          pure (TH.VarE n, TH.LamE [TH.VarP n])
    prep asg af = do
      (as, g) <- asg
      (a, f) <- af
      pure (TH.AppE (TH.VarE 'displayish) a : as, f . g)

interpolator ::
  Char ->
  -- | postprocess the 'TH.Exp'
  (TH.Exp -> TH.Exp) ->
  TH.QuasiQuoter
interpolator c pp = TH.QuasiQuoter {..}
  where
    quoteExp = parseSegments c >=> ipExpr pp
    quotePat = const $ fail "pattern context not supported"
    quoteType = const $ fail "type context not supported"
    quoteDec = const $ fail "declaration context not supported"

class Displayish a where
  displayish :: a -> TLB.Builder

instance {-# OVERLAPPABLE #-} TD.Display a => Displayish a where
  displayish = TD.displayBuilder

-- String is still used too pervasively...
instance Displayish String where
  displayish = TLB.fromString

class Stringish a where
  stringish :: TLB.Builder -> a

instance Stringish String where
  stringish = TL.unpack . TLB.toLazyText

instance Stringish T.Text where
  stringish = TL.toStrict . TLB.toLazyText

instance Stringish TL.Text where
  stringish = TLB.toLazyText