{-# LANGUAGE OverloadedStrings #-}
{-# LANGUAGE Safe #-}
module Embedder.Atmo.Util
      ( FieldId
      , paramTys, isPrim, inlineable, mustInline, nil
      , flattenTuple, mkTuple, mkTuplePat
      , mkPair, mkPairPat, flattenLam, mkLam
      , mkApp, mkError, builtin, proxy, tybuiltin, userBuiltin, mkTupleCtor, isTupleCtor
      , mkRecVal, mkRecUpd, mkRecSel, mkPatRec
      ) where

import ReWire.Annotation (Annote (MsgAnnote))
import Embedder.Atmo.Syntax (Exp (..), Ty (..), Pat (..), Defn (..), DefnAttr (..), Poly (..), FieldId, TyBuiltin)
import Embedder.Atmo.Types (proxyTy, nilTy, strTy, arr, typeOf, arrowRight, pairTy, fundamental, mkArrowTy, paramTys)
import Embedder.Builtins ( builtins, Builtin (..), tybuiltins, builtinUserQName, RWUserOp (..) )

import Data.Text (Text, unpack)
import Data.Maybe(fromMaybe)
import Numeric.Natural (Natural)
import Data.Tuple (swap)
import qualified Data.Text as T


-- Name Handling:
      -- Remove Embed
      -- Remove bindings (e.g. Poly)
      -- Remove Name from Lam helpers
      -- Remove s2n and n2s
      -- Remove Unbound exports

builtin :: Text -> Maybe Builtin
builtin :: Text -> Maybe Builtin
builtin Text
b = Text -> [(Text, Builtin)] -> Maybe Builtin
forall a b. Eq a => a -> [(a, b)] -> Maybe b
lookup Text
b [(Text, Builtin)]
builtins

tybuiltin :: Text -> Maybe TyBuiltin
tybuiltin :: Text -> Maybe TyBuiltin
tybuiltin Text
b = Text -> [(Text, TyBuiltin)] -> Maybe TyBuiltin
forall a b. Eq a => a -> [(a, b)] -> Maybe b
lookup Text
b [(Text, TyBuiltin)]
tybuiltins

userBuiltin :: Text -> Maybe Builtin
userBuiltin :: Text -> Maybe Builtin
userBuiltin Text
b = Text -> [(Text, Builtin)] -> Maybe Builtin
forall a b. Eq a => a -> [(a, b)] -> Maybe b
lookup Text
b (((Builtin, Text) -> (Text, Builtin))
-> [(Builtin, Text)] -> [(Text, Builtin)]
forall a b. (a -> b) -> [a] -> [b]
map (Builtin, Text) -> (Text, Builtin)
forall a b. (a, b) -> (b, a)
swap [(Builtin, Text)]
builtinUserQName)

isPrim :: Show a => a -> Bool
isPrim :: forall a. Show a => a -> Bool
isPrim = Char -> [Char] -> Bool
forall (t :: * -> *) a. (Foldable t, Eq a) => a -> t a -> Bool
notElem Char
'.' ([Char] -> Bool) -> (a -> [Char]) -> a -> Bool
forall b c a. (b -> c) -> (a -> b) -> a -> c
. a -> [Char]
forall a. Show a => a -> [Char]
show

inlineable :: Defn -> Bool
inlineable :: Defn -> Bool
inlineable Defn
d = case Defn -> Maybe DefnAttr
defnAttr Defn
d of
      Just DefnAttr
Inline   -> Bool
True
      Just DefnAttr
NoInline -> Bool
False
      Maybe DefnAttr
Nothing       -> Bool -> Bool
not (Bool -> Bool) -> Bool -> Bool
forall a b. (a -> b) -> a -> b
$ Text -> Bool
forall a. Show a => a -> Bool
isPrim (Text -> Bool) -> Text -> Bool
forall a b. (a -> b) -> a -> b
$ Defn -> Text
defnName Defn
d

mustInline :: Defn -> Bool
mustInline :: Defn -> Bool
mustInline = \ case
      Defn { defnAttr :: Defn -> Maybe DefnAttr
defnAttr = Just DefnAttr
Inline }                    -> Bool
True
      Defn { defnPolyTy :: Defn -> Poly
defnPolyTy = Poly [Text]
_ Ty
t
           , defnName :: Defn -> Text
defnName   = Text
n }                            -> Bool -> Bool
not (Text -> Bool
forall a. Show a => a -> Bool
isPrim Text
n Bool -> Bool -> Bool
|| Ty -> Bool
fundamental Ty
t)

nil :: Exp
nil :: Exp
nil = Annote -> Maybe Poly -> Maybe Ty -> Text -> Exp
Con (Text -> Annote
MsgAnnote Text
"nil") Maybe Poly
forall a. Maybe a
Nothing (Ty -> Maybe Ty
forall a. a -> Maybe a
Just Ty
nilTy) Text
"()"

proxy :: Natural -> Exp
proxy :: Natural -> Exp
proxy Natural
n = Annote -> Maybe Poly -> Maybe Ty -> Text -> Exp
Con Annote
an Maybe Poly
forall a. Maybe a
Nothing (Ty -> Maybe Ty
forall a. a -> Maybe a
Just (Ty -> Maybe Ty) -> Ty -> Maybe Ty
forall a b. (a -> b) -> a -> b
$ Annote -> Natural -> Ty
proxyTy Annote
an Natural
n) Text
"Proxy"
      where an :: Annote
            an :: Annote
an = Text -> Annote
MsgAnnote Text
"Proxy"

nilPat :: Pat
nilPat :: Pat
nilPat = Annote -> Maybe Poly -> Maybe Ty -> Text -> [Pat] -> Pat
PatCon (Text -> Annote
MsgAnnote Text
"nilPat") Maybe Poly
forall a. Maybe a
Nothing (Ty -> Maybe Ty
forall a. a -> Maybe a
Just Ty
nilTy) Text
"()" []

mkPair :: Annote -> Exp -> Exp -> Exp
mkPair :: Annote -> Exp -> Exp -> Exp
mkPair Annote
an Exp
e1 Exp
e2 = Annote -> Exp -> [Exp] -> Exp
mkApp Annote
an (Annote -> Maybe Poly -> Maybe Ty -> Text -> Exp
Con Annote
an Maybe Poly
forall a. Maybe a
Nothing Maybe Ty
t Text
"(,)") [Exp
e1, Exp
e2]
      where t :: Maybe Ty
            t :: Maybe Ty
t = do
                  t1 <- Exp -> Maybe Ty
forall a. TypeAnnotated a => a -> Maybe Ty
typeOf Exp
e1
                  t2 <- typeOf e2
                  pure $ mkArrowTy [t1, t2] $ pairTy an t1 t2

flattenPair :: Exp -> Maybe (Exp,Exp)
flattenPair :: Exp -> Maybe (Exp, Exp)
flattenPair = \ case
      Tuple Annote
_ Maybe Poly
_ Maybe Ty
_ [Exp
e1,Exp
e2] -> (Exp, Exp) -> Maybe (Exp, Exp)
forall a. a -> Maybe a
Just (Exp
e1,Exp
e2)
      Exp
_ -> Maybe (Exp, Exp)
forall a. Maybe a
Nothing

mkPairPat :: Annote -> Pat -> Pat -> Pat
mkPairPat :: Annote -> Pat -> Pat -> Pat
mkPairPat Annote
an Pat
p1 Pat
p2 = Annote -> Maybe Poly -> Maybe Ty -> Text -> [Pat] -> Pat
PatCon Annote
an Maybe Poly
forall a. Maybe a
Nothing (Annote -> Ty -> Ty -> Ty
pairTy Annote
an (Ty -> Ty -> Ty) -> Maybe Ty -> Maybe (Ty -> Ty)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Pat -> Maybe Ty
forall a. TypeAnnotated a => a -> Maybe Ty
typeOf Pat
p1 Maybe (Ty -> Ty) -> Maybe Ty -> Maybe Ty
forall a b. Maybe (a -> b) -> Maybe a -> Maybe b
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> Pat -> Maybe Ty
forall a. TypeAnnotated a => a -> Maybe Ty
typeOf Pat
p2) Text
"(,)" [Pat
p1, Pat
p2]

mkTupleCtor :: Int -> Text
mkTupleCtor :: Int -> Text
mkTupleCtor Int
n = Text
"(" Text -> Text -> Text
forall a. Semigroup a => a -> a -> a
<> Int -> Text -> Text
T.replicate (Int
nInt -> Int -> Int
forall a. Num a => a -> a -> a
-Int
1) Text
"," Text -> Text -> Text
forall a. Semigroup a => a -> a -> a
<> Text
")"

-- | Function to match the specified pattern and return Maybe Int
isTupleCtor :: T.Text -> Bool
isTupleCtor :: Text -> Bool
isTupleCtor Text
c = Text
c Text -> Text -> Bool
forall a. Eq a => a -> a -> Bool
== Text
"(" Text -> Text -> Text
forall a. Semigroup a => a -> a -> a
<> Int -> Text -> Text
T.replicate ([Char] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length (Text -> [Char]
unpack Text
c) Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
2) Text
"," Text -> Text -> Text
forall a. Semigroup a => a -> a -> a
<> Text
")"

mkTuple :: Annote -> [Exp] -> Exp
mkTuple :: Annote -> [Exp] -> Exp
mkTuple Annote
an = (Exp -> Exp -> Exp) -> Exp -> [Exp] -> Exp
forall a b. (a -> b -> b) -> b -> [a] -> b
forall (t :: * -> *) a b.
Foldable t =>
(a -> b -> b) -> b -> t a -> b
foldr (Annote -> Exp -> Exp -> Exp
mkPair Annote
an) Exp
nil

flattenTuple :: Exp -> Maybe [Exp]
flattenTuple :: Exp -> Maybe [Exp]
flattenTuple Exp
e = case Exp -> Maybe (Exp, Exp)
flattenPair Exp
e of
      Maybe (Exp, Exp)
Nothing -> Maybe [Exp]
forall a. Maybe a
Nothing
      -- Can't remove unit without changing type. Just (e1,Con _ _ _ (n2s -> "()")) -> Just [e1]
      Just (Exp
e1,Exp
e2) -> [Exp] -> Maybe [Exp]
forall a. a -> Maybe a
Just ([Exp] -> Maybe [Exp]) -> [Exp] -> Maybe [Exp]
forall a b. (a -> b) -> a -> b
$ Exp
e1 Exp -> [Exp] -> [Exp]
forall a. a -> [a] -> [a]
: [Exp] -> Maybe [Exp] -> [Exp]
forall a. a -> Maybe a -> a
fromMaybe [Exp
e2] (Exp -> Maybe [Exp]
flattenTuple Exp
e2)

mkTuplePat :: Annote -> [Pat] -> Pat
mkTuplePat :: Annote -> [Pat] -> Pat
mkTuplePat Annote
an = (Pat -> Pat -> Pat) -> Pat -> [Pat] -> Pat
forall a b. (a -> b -> b) -> b -> [a] -> b
forall (t :: * -> *) a b.
Foldable t =>
(a -> b -> b) -> b -> t a -> b
foldr (Annote -> Pat -> Pat -> Pat
mkPairPat Annote
an) Pat
nilPat

mkApp :: Annote -> Exp -> [Exp] -> Exp
mkApp :: Annote -> Exp -> [Exp] -> Exp
mkApp Annote
an Exp
e [Exp]
es = Annote -> Maybe Poly -> Maybe Ty -> Exp -> [Exp] -> Exp
App Annote
an Maybe Poly
forall a. Maybe a
Nothing ((Maybe Ty -> Exp -> Maybe Ty) -> Maybe Ty -> [Exp] -> Maybe Ty
forall b a. (b -> a -> b) -> b -> [a] -> b
forall (t :: * -> *) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
foldl' (\ Maybe Ty
b Exp
_ -> (Ty -> Ty) -> Maybe Ty -> Maybe Ty
forall a b. (a -> b) -> Maybe a -> Maybe b
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap Ty -> Ty
arrowRight Maybe Ty
b) (Exp -> Maybe Ty
forall a. TypeAnnotated a => a -> Maybe Ty
typeOf Exp
e) [Exp]
es) Exp
e [Exp]
es

mkError :: Annote -> Maybe Ty -> Text -> Exp
mkError :: Annote -> Maybe Ty -> Text -> Exp
mkError Annote
an Maybe Ty
t Text
err = Annote -> Maybe Poly -> Maybe Ty -> Exp -> [Exp] -> Exp
App Annote
an Maybe Poly
forall a. Maybe a
Nothing Maybe Ty
t (Annote -> Maybe Poly -> Maybe Ty -> RWUserOp -> Exp
RWUser Annote
an Maybe Poly
forall a. Maybe a
Nothing (Ty -> Ty -> Ty
arr (Annote -> Ty
strTy Annote
an) (Ty -> Ty) -> Maybe Ty -> Maybe Ty
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Maybe Ty
t) (Builtin -> RWUserOp
RWBuiltin Builtin
Error)) [Annote -> Maybe Poly -> Text -> Exp
LitStr Annote
an Maybe Poly
forall a. Maybe a
Nothing Text
err]

flattenLam :: Monad m => Exp -> m ([Text], Exp)
flattenLam :: forall (m :: * -> *). Monad m => Exp -> m ([Text], Exp)
flattenLam = \ case
      Lam Annote
_ Maybe Poly
_ Maybe Ty
_ [Text]
xs Exp
e -> do
            (xs', e') <- Exp -> m ([Text], Exp)
forall (m :: * -> *). Monad m => Exp -> m ([Text], Exp)
flattenLam Exp
e
            pure (xs ++ xs', e')
      Exp
e              -> ([Text], Exp) -> m ([Text], Exp)
forall a. a -> m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure ([], Exp
e)

mkLam :: Annote -> Ty -> [Text] -> Exp -> Exp
mkLam :: Annote -> Ty -> [Text] -> Exp -> Exp
mkLam Annote
an Ty
t = Annote -> Maybe Poly -> Maybe Ty -> [Text] -> Exp -> Exp
Lam Annote
an Maybe Poly
forall a. Maybe a
Nothing (Ty -> Maybe Ty
forall a. a -> Maybe a
Just Ty
t)



-- Record helpers for construction
mkRecVal :: Annote -> [(Text, Exp)] -> Exp
mkRecVal :: Annote -> [(Text, Exp)] -> Exp
mkRecVal Annote
an = Annote -> Maybe Poly -> Maybe Ty -> [(Text, Exp)] -> Exp
RecVal Annote
an Maybe Poly
forall a. Maybe a
Nothing Maybe Ty
forall a. Maybe a
Nothing

mkRecUpd :: Annote -> Exp -> [(Text, Exp)] -> Exp
mkRecUpd :: Annote -> Exp -> [(Text, Exp)] -> Exp
mkRecUpd Annote
an = Annote -> Maybe Poly -> Maybe Ty -> Exp -> [(Text, Exp)] -> Exp
RecUpd Annote
an Maybe Poly
forall a. Maybe a
Nothing Maybe Ty
forall a. Maybe a
Nothing

mkRecSel :: Annote -> Text -> Exp -> Exp
mkRecSel :: Annote -> Text -> Exp -> Exp
mkRecSel Annote
an = Annote -> Maybe Poly -> Maybe Ty -> Text -> Exp -> Exp
RecSel Annote
an Maybe Poly
forall a. Maybe a
Nothing Maybe Ty
forall a. Maybe a
Nothing

-- Record pattern helper
mkPatRec :: Annote -> [(Text, Pat)] -> Pat
mkPatRec :: Annote -> [(Text, Pat)] -> Pat
mkPatRec Annote
an = Annote -> Maybe Poly -> Maybe Ty -> [(Text, Pat)] -> Pat
PatRec Annote
an Maybe Poly
forall a. Maybe a
Nothing Maybe Ty
forall a. Maybe a
Nothing