Weld.hs

module Weld where

import Control.Arrow

import Data.Bifunctor as B

import Text.PrettyPrint as PP hiding ((<>))

infix 1 \/
type (\/) = Either

data Prog a b where
  ID      :: Prog a a
  (:#)    :: Prog a b -> Prog b c -> Prog a c

  (:&)    :: Prog a b -> Prog a c -> Prog a (b, c)
  FST     :: Prog (a, b) a
  SND     :: Prog (a, b) b

  IT      :: Prog a ()

  APPLY   :: Prog (Prog a b, a) b
  CURRY   :: Prog (a, b) c -> Prog a (Prog b c)
  UNCURRY :: Prog a (Prog b c) -> Prog (a, b) c

  CONST   :: a -> Prog () a

  (:|)    :: Prog a c -> Prog b c -> Prog (a \/ b)  c
  LEFT    :: Prog a (a \/ b)
  RIGHT   :: Prog b (a \/ b)

  CAST    :: Same a b => Prog a b

  EXT     :: String -> (a -> b) -> Prog a b
  FIX     :: Prog (Prog a b, a) b -> Prog a b
  DUP     :: Prog a (a, a)
  ALWAYS  :: Prog a (Prog () a)
  CONS    :: Prog (Prog a b, Prog b c) (Prog a c)


  DISTR   :: Prog (a, b \/ c) ((a, b) \/ (a, c))
  GATHER  :: Prog (a \/ c, b \/ c) ((a, b) \/ c)

  deriving Show via PP (Prog a b)

class Same a b | a -> b, b -> a where
  forth :: a -> b

instance Same [a] ((a, [a]) \/ ()) where
  forth [] = Right ()
  forth (x : xs) = Left (x, xs)

instance Same ((a, [a]) \/ ()) [a] where
  forth = either (uncurry (:)) (const [])

class Pretty a where
  pp :: a -> Doc

newtype PP a = PP { unPP :: a }

instance Pretty a => Show (PP a) where
  show = show . pp . unPP

instance Pretty (Prog a b) where
  pp t = case t of
    CONST _ -> "<smth>"
    CAST -> "CAST"
    EXT n _ -> text n
    APPLY -> "APPLY"
    UNCURRY f -> "UNCURRY" `indent` parens (pp f)
    ID -> "."
    FIX p -> "FIX" `indent` parens (pp p)
    DUP -> "DUP"
    FST -> "FST"
    SND -> "SND"
    LEFT -> "LEFT"
    RIGHT -> "RIGHT"
    _ :# _ -> collectHash t
    _ :| _ -> collectOr   t
    _ :& _ -> collectAnd  t
    ALWAYS -> "ALWAYS"
    CONS -> "CONS"
    DISTR -> "DISTR"
    GATHER -> "GATHER"
    IT -> "IT"
    CURRY f -> "CURRY" `indent` parens (pp f)

collectHash :: Prog a b -> Doc
collectHash = \case
  a :# b -> collectHash a PP.<+> "#" PP.<+> collectHash b
  x      -> pp x

collectOr :: Prog a b -> Doc
collectOr t = case t of
  a :# b -> parens (pp a) PP.<+> "|" PP.<+> parens (pp b)
  a :| b -> parens (pp a) PP.<+> "|" PP.<+> parens (pp b)
  _      -> pp t

collectAnd :: Prog a b -> Doc
collectAnd t = case t of
  a :# b -> parens (pp a) PP.<+> "&" PP.<+> parens (pp b)
  a :| b -> parens (pp a) PP.<+> "&" PP.<+> parens (pp b)
  a :& b -> parens (pp a) PP.<+> "&" PP.<+> parens (pp b)
  _      -> pp t

indent, above :: Doc -> Doc -> Doc
indent a b = hang a 2 b
above  a b = hang a 0 b

aSSOC :: Prog (a, (b, c)) ((a, b), c)
aSSOC = (FST :& (SND :# FST)) :& (SND :# SND)

uNASSOC :: Prog ((a, b), c) (a, (b, c))
uNASSOC = (FST :# FST) :& ((FST :# SND) :& SND)

sWAP :: Prog (a, b) (b, a)
sWAP = SND :& FST

(%&) :: Prog a b -> Prog c d -> Prog (a, c) (b, d)
f %& g = (FST :# f) :& (SND :# g)

(%|) :: Prog a b -> Prog c d -> Prog (a \/ c) (b \/ d)
f %| g = (f :# LEFT) :| (g :# RIGHT)

distribute :: Bifunctor d => a -> d b c -> d (a, b) (a, c)
distribute a = bimap (a,) (a,)

class    Swap  d   where swap :: d a b -> d b a
instance Swap (,)  where swap (a, b) = (b, a)
instance Swap (\/) where swap = either Right Left

class    Gather  d   where gather :: d (a \/ c) (b \/ c) -> d a b \/ c
instance Gather (,)  where gather (ma, mb) = either (\a -> B.first (a,) mb) Right ma
instance Gather (\/) where gather = either (either (Left . Left)  Right) (either (Left . Right) Right)

class Assoc d where
  assoc   :: d a (d b c) -> d (d a b) c
  unassoc :: d (d a b) c -> d a (d b c)

instance Assoc (,) where
  assoc   (a, (b, c)) = ((a, b), c)
  unassoc ((a, b), c) = (a, (b, c))

instance Assoc (\/) where
  assoc   = either (Left . Left) (either (Left . Right) Right)
  unassoc = either (either Left (Right . Left)) (Right . Right)

rIGHT :: Prog (a, b \/ c) (b \/ (a, c))
rIGHT = DISTR :# (SND %| ID)

lEFT :: Prog (a, b \/ c) ((a, b) \/ c)
lEFT = DISTR :# (ID %| SND)

aNY :: Prog (a \/ a) a
aNY = ID :| ID

run :: Prog a b -> a -> b
run = \case
  CONST a -> const a
  CAST    -> forth

  EXT _ f -> f

  ID      -> id
  FIX   f -> \a -> run f (FIX f, a)
  APPLY   -> uncurry run
  UNCURRY f -> \(a, b) -> run (run f a) b
  CURRY   f -> \a -> _

  DUP     -> \a -> (a, a)
  FST     -> fst
  SND     -> snd

  LEFT    -> Left
  RIGHT   -> Right

  f :# g  -> run f >>> run g
  f :& g  -> run f &&& run g
  f :| g  -> run f ||| run g

  ALWAYS -> CONST
  CONS   -> uncurry (:#)

  DISTR  -> uncurry distribute
  GATHER -> gather

-- fix :: (a -> a) -> a
-- fix f = x where x = f x

-- -- foldList :: forall a b. Prog (Prog (a, b) b, (b, [a])) b
-- -- foldList =
-- --   LET \f ->
-- --   LET \b ->
-- --   FIX do
-- --     LET \self ->
-- --       CAST @[a]
-- --       :# (ID :% self :# f) :% CONST b
-- --       :# aNY

-- -- mapList :: forall a b. Prog (Prog a b, [a]) [b]
-- -- mapList =
-- --   LET \f -> DUP :# (DUP :# (_ :% ID) :# (CONST [] :% ID)) :# _ :# foldList