Weird little prefix catlang

Yawn.hs

{-# Language LambdaCase #-}

import Control.Applicative (Alternative (empty), liftA2)
import Control.Category ((>>>))
import Control.Monad (guard, join, void)
import Control.Monad.State.Strict
  (
    StateT (StateT, runStateT),
    gets,
    modify,
  )
import Control.Monad.Trans.Writer.CPS
  (
    WriterT,
    mapWriterT,
    runWriterT,
    tell,
    writerT,
  )
import Control.Monad.Trans.Class (lift)
import Data.Char (isAsciiLower, isAsciiUpper, isDigit)
import Data.Foldable (traverse_)
import Data.Map.Strict (Map)
import Data.Map.Strict qualified as Map
import Data.Maybe (fromMaybe)
import Text.ParserCombinators.ReadP
  (
    (<++),
    ReadP,
    between,
    eof,
    look,
    munch,
    munch1,
    readP_to_S,
    readS_to_P,
    satisfy,
    skipMany,
    skipSpaces,
    string,
  )

type P =
  StateT (Map String String)
    (WriterT [String]
      ReadP)

prelude :: String
prelude =
  unlines [
    "__define__ def (__define__)",
    "def note (__drop__)",
    "def disabled (__drop__)",
    "",
    "def pair   (cat unit dip(unit))",
    "def unpair (i)",
    "",
    "def force   (forcing())",
    "def forcing (unpair i dip(pair))",
    "",
    "def zap  (__drop__)",
    "def i    (__apply__)",
    "def unit (__quote__)",
    "def rep  (i cat dup)",
    "def m    (i dup)",
    "def run  (i cat unit dup)",
    "def dup  (__dup__)",
    "def k    (i zap swap)",
    "def z    (i zap)",
    "def nip  (zap swap)",
    "def sap  (i cat swap)",
    "def t    (i swap)",
    "def dip  (i cat unit swap)",
    "def cat  (__compose__)",
    "def swat (cat swap)",
    "def swap (__swap__)",
    "def cons (cat swap unit swap)",
    "def take (cat unit swap)",
    "def tack (cat unit)",
    "def sip  (dip dip(dup))",
    "def w    (i force dip(dup))",
    "def peek (swap dip(dup))",
    "disabled (def poke (swap dip(dip(zap))))",
    "def poke (swap unpair zap swap pair)",
    "def cake (",
    "    force dip(cat swap)",
    "    swap",
    "    dip(cat)",
    "    dup",
    "    dip(dup unit)",
    ")",
    "def b (i force dip(cons))",
    "def c (i force dip(swap))",
    "def dig (force dip(unpair) swap pair)",
    "def bury (force dip(swap) swap)",
    "def flip (swap unpair swap dip(pair))",
    "def s (i force dip(swap cons dip(dup)))",
    ""
  ]

main :: IO ()
main = do

  check prelude "zap  (A)"                 ""
  check prelude "i    (A)"                 "A"
  check prelude "unit (A)"                 "((A))"
  check prelude "rep  (A)"                 "A A"
  check prelude "m    (A)"                 "A (A)"
  check prelude "run  (A)"                 "(A) A"
  check prelude "dup  (A)"                 "(A) (A)"
  check prelude "k    (A) (B)"             "A"
  check prelude "z    (A) (B)"             "B"
  check prelude "nip  (A) (B)"             "(A)"
  check prelude "sap  (A) (B)"             "B A"
  check prelude "t    (A) (B)"             "B (A)"
  check prelude "dip  (A) (B)"             "(B) A"
  check prelude "cat  (A) (B)"             "(A B)"
  check prelude "swat (A) (B)"             "(B A)"
  check prelude "swap (A) (B)"             "(B) (A)"
  check prelude "cons (A) (B)"             "(A (B))"
  check prelude "take (A) (B)"             "((B) A)"
  check prelude "tack (A) (B)"             "((A) B)"
  check prelude "sip  (A) (B)"             "(B) A (B)"
  check prelude "w    (A) (B)"             "A (B) (B)"
  check prelude "peek (A) (B)"             "(B) (A) (B)"
  check prelude "cake (A) (B)"             "(A (B)) ((B) A)"
  check prelude "poke (A) (B) (C)"         "(B) (A)"
  check prelude "b    (A) (B) (C)"         "A (B (C))"
  check prelude "c    (A) (B) (C)"         "A (C) (B)"
  check prelude "dig  (A) (B) (C)"         "(C) (A) (B)"
  check prelude "bury (A) (B) (C)"         "(B) (C) (A)"
  check prelude "flip (A) (B) (C)"         "(C) (B) (A)"
  check prelude "s    (A) (B) (C)"         "A (C) (B (C))"
  check prelude "s'   (A) (B) (C) (D)"     "B (D) A (C (D))"
  check prelude "j    (A) (B) (C) (D)"     "A (B) (A (D) (C))"
  check prelude "j'   (A) (B) (C) (D) (E)" "B (C) (B (E) A (D))"

  pure ()

check :: String -> String -> String -> IO ()
check prelude source expected =
  either failed passed
    (test (prelude <> source) expected)
  where
    failed actual =
      (putStrLn . unlines) [
        unwords ["failed:", show source],
        unwords ["actual:", show actual],
        unwords ["expected:", show expected]
      ]
    passed actual =
      (putStrLn . unlines) [
        unwords ["passed:", show source, "=>", show actual]
      ]

test :: String -> String -> Either String String
test source expected =
  maybe
    (Left "")
    (\(_logs, actual) ->
      if actual == expected
        then Right actual
        else Left actual)
    (run source)

trace :: String -> IO ()
trace =
  traverse_ log .
  maybe [] (\(logs, end) -> logs <> [end]) .
  run
  where
    log state = do
      putStrLn state
      putStrLn (replicate 32 '-')

run :: String -> Maybe ([String], String)
run =
  one .
  fmap (\((((), _defs), logs), end) -> (logs, end)) .
  readP_to_S
    (skipSpaces *> runWriterT (runStateT steps mempty))

steps :: P ()
steps = void (greedily step)

step :: P ()
step = do
  lift . tell . (: []) =<< (lift . lift) look
  first
    [
      step_dup,
      step_swap,
      step_drop,
      step_apply,
      step_quote,
      step_compose,
      step_define,
      step_use
    ]

step_dup :: P ()
step_dup = do
  keyword "__dup__"
  a <- get
  put a
  put a

step_swap :: P ()
step_swap = do
  keyword "__swap__"
  a <- get
  b <- get
  put a
  put b

step_drop :: P ()
step_drop = do
  keyword "__drop__"
  void get

step_apply :: P ()
step_apply = do
  keyword "__apply__"
  put =<< open

step_quote :: P ()
step_quote = do
  keyword "__quote__"
  put . wrap =<< get

step_compose :: P ()
step_compose = do
  keyword "__compose__"
  a <- open
  b <- open
  put (wrap (a `beside` b))

step_define :: P ()
step_define = do
  keyword "__define__"
  x <- word
  a <- wrapped
  modify (Map.insert x a)

step_use :: P ()
step_use = do
  x <- word
  maybe empty put =<< gets (Map.lookup x)

keyword :: String -> P ()
keyword = (word >>=) . same

same :: (Alternative f, Eq a) => a -> a -> f ()
same = fmap guard . (==)

get :: P String
get = alt (step *> get) term

term :: P String
term = first [block, word]

block :: P String
block = fmap wrap wrapped

wrap :: String -> String
wrap = ("(" <>) >>> (<> ")")

wrapped :: P String
wrapped = blocked (fmap unwords (greedily term))

blocked :: P a -> P a
blocked = (symbol "(" *>) >>> (<* symbol ")")

symbol :: String -> P String
symbol = token . lift . lift . string

token :: P a -> P a
token = (<* (lift . lift) skipSpaces)

word :: P String
word = token
  ((lift . lift) (name <++ punctuation <++ operator))
  where
    name = liftA2 (:)
      (satisfy beginsName)
      (munch isName)
    beginsName =
      isAsciiLower <||>
      isAsciiUpper <||>
      (`elem` "'_")
    isName =
      beginsName <||>
      isDigit <||>
      (`elem` "-")
    punctuation =
      fmap pure (satisfy isPunctuation)
    isPunctuation =
      (`elem` ",;")
    operator =
      munch1 isOperator
    isOperator =
      (`elem` "!#$%&*+-./:<=>?@\\^|~")

(<||>) :: (Applicative f) => f Bool -> f Bool -> f Bool
(<||>) = liftA2 (||)
infixr 2 <||>

put :: String -> P ()
put = next . beside

open :: P String
open = unwrap =<< get

unwrap :: String -> P String
unwrap s = StateT \defs ->
  (
    writerT .
    one .
    fmap fst .
    readP_to_S (runWriterT (runStateT wrapped defs) <* eof)
  ) s

beside :: String -> String -> String
beside = curry \case
  ("", b) -> b
  (a, "") -> a
  (a, b) -> unwords [a, b]

first :: [P a] -> P a
first = foldr alt empty

greedily :: P a -> P [a]
greedily p = loop
  where
    loop = liftA2 (:) p loop `alt` pure []

alt :: P a -> P a -> P a
alt p1 p2 = StateT \defs ->
  writerT
    (runWriterT (runStateT p1 defs) <++
      runWriterT (runStateT p2 defs))

rest :: P String
rest = (lift . lift) look <* replace mempty

replace :: String -> P ()
replace = next . const

next :: (String -> String) -> P ()
next = lift . lift . readS_to_P . fmap (pure . (,) ())

one :: (Alternative f) => [a] -> f a
one = \case
  [s] -> pure s
  ss -> empty