{-# LANGUAGE OverloadedStrings #-}

-- |
-- SPDX-License-Identifier: BSD-3-Clause
--
-- Term elaboration which happens after type checking.
module Swarm.Language.Elaborate where

import Control.Applicative ((<|>))
import Control.Lens (transform, (^.))
import Swarm.Language.Syntax

-- | Perform some elaboration / rewriting on a fully type-annotated
--   term.  This currently performs such operations as rewriting @if@
--   expressions and recursive let expressions to use laziness
--   appropriately.  It also inserts types and requirements on bound
--   variables so they will be available at runtime.
--
--   In theory it could also perform rewriting for overloaded
--   constants depending on the actual type they are used at, but
--   currently that sort of thing tends to make type inference fall
--   over.
elaborate :: TSyntax -> TSyntax
elaborate :: TSyntax -> TSyntax
elaborate = (TSyntax -> TSyntax) -> TSyntax -> TSyntax
forall a. Plated a => (a -> a) -> a -> a
transform TSyntax -> TSyntax
rewrite
 where
  rewrite :: TSyntax -> TSyntax
  rewrite :: TSyntax -> TSyntax
rewrite (Syntax' SrcLoc
l Term' (Poly Type)
t Comments
cs Poly Type
ty) = SrcLoc -> Term' (Poly Type) -> Comments -> Poly Type -> TSyntax
forall ty. SrcLoc -> Term' ty -> Comments -> ty -> Syntax' ty
Syntax' SrcLoc
l (Term' (Poly Type) -> Term' (Poly Type)
rewriteTerm Term' (Poly Type)
t) Comments
cs Poly Type
ty

  rewriteTerm :: TTerm -> TTerm
  rewriteTerm :: Term' (Poly Type) -> Term' (Poly Type)
rewriteTerm = \case
    -- Here we take inferred types for variables bound by def or
    -- bind (but NOT let) and stuff them into the term itself, so that
    -- we will still have access to them at runtime, after type
    -- annotations on the AST are erased.  We need them at runtime so
    -- we can keep track of the types of variables in scope, for use
    -- in typechecking additional terms entered at the REPL.  The
    -- reason we do not do this for 'let' is so that 'let' introduces
    -- truly local bindings which will not be available for use in
    -- later REPL terms.
    --
    -- We assume requirements for these variables have already been
    -- filled in during typechecking.  The reason we need to wait
    -- until now to fill in the types is that we have to wait until
    -- solving, substitution, and generalization are complete.
    --
    -- Eventually, once requirements analysis is part of typechecking,
    -- we'll infer them both at typechecking time then fill them in
    -- during elaboration here.
    SLet LetSyntax
ls Bool
r LocVar
x Maybe (Poly Type)
mty Maybe Requirements
mreq TSyntax
t1 TSyntax
t2 ->
      let mty' :: Maybe (Poly Type)
mty' = case LetSyntax
ls of
            LetSyntax
LSDef -> Maybe (Poly Type)
mty Maybe (Poly Type) -> Maybe (Poly Type) -> Maybe (Poly Type)
forall a. Maybe a -> Maybe a -> Maybe a
forall (f :: * -> *) a. Alternative f => f a -> f a -> f a
<|> Poly Type -> Maybe (Poly Type)
forall a. a -> Maybe a
Just (TSyntax
t1 TSyntax -> Getting (Poly Type) TSyntax (Poly Type) -> Poly Type
forall s a. s -> Getting a s a -> a
^. Getting (Poly Type) TSyntax (Poly Type)
forall ty (f :: * -> *).
Functor f =>
(ty -> f ty) -> Syntax' ty -> f (Syntax' ty)
sType)
            LetSyntax
LSLet -> Maybe (Poly Type)
forall a. Maybe a
Nothing
       in LetSyntax
-> Bool
-> LocVar
-> Maybe (Poly Type)
-> Maybe Requirements
-> TSyntax
-> TSyntax
-> Term' (Poly Type)
forall ty.
LetSyntax
-> Bool
-> LocVar
-> Maybe (Poly Type)
-> Maybe Requirements
-> Syntax' ty
-> Syntax' ty
-> Term' ty
SLet LetSyntax
ls Bool
r LocVar
x Maybe (Poly Type)
mty' Maybe Requirements
mreq TSyntax
t1 TSyntax
t2
    SBind Maybe LocVar
x (Just Poly Type
ty) Maybe (Poly Type)
_ Maybe Requirements
mreq TSyntax
c1 TSyntax
c2 -> Maybe LocVar
-> Maybe (Poly Type)
-> Maybe (Poly Type)
-> Maybe Requirements
-> TSyntax
-> TSyntax
-> Term' (Poly Type)
forall ty.
Maybe LocVar
-> Maybe ty
-> Maybe (Poly Type)
-> Maybe Requirements
-> Syntax' ty
-> Syntax' ty
-> Term' ty
SBind Maybe LocVar
x Maybe (Poly Type)
forall a. Maybe a
Nothing (Poly Type -> Maybe (Poly Type)
forall a. a -> Maybe a
Just Poly Type
ty) Maybe Requirements
mreq TSyntax
c1 TSyntax
c2
    -- Rewrite @f $ x@ to @f x@.
    SApp (Syntax' SrcLoc
_ (SApp (Syntax' SrcLoc
_ (TConst Const
AppF) Comments
_ Poly Type
_) TSyntax
l) Comments
_ Poly Type
_) TSyntax
r -> TSyntax -> TSyntax -> Term' (Poly Type)
forall ty. Syntax' ty -> Syntax' ty -> Term' ty
SApp TSyntax
l TSyntax
r
    -- Leave any other subterms alone.
    Term' (Poly Type)
t -> Term' (Poly Type)
t

-- | Insert a special 'suspend' primitive at the very end of an erased
--   term which must have a command type.
insertSuspend :: Term -> Term
insertSuspend :: Term -> Term
insertSuspend Term
t = case Term
t of
  -- Primitive things which have type Cmd Unit: p => (p ; suspend ())
  TRequireDevice {} -> Term
thenSuspend
  TRequire {} -> Term
thenSuspend
  TRequirements {} -> Term
thenSuspend
  -- Recurse through def, tydef, bind, and annotate.
  TLet LetSyntax
ls Bool
r Var
x Maybe (Poly Type)
mty Maybe Requirements
mreq Term
t1 Term
t2 -> LetSyntax
-> Bool
-> Var
-> Maybe (Poly Type)
-> Maybe Requirements
-> Term
-> Term
-> Term
TLet LetSyntax
ls Bool
r Var
x Maybe (Poly Type)
mty Maybe Requirements
mreq Term
t1 (Term -> Term
insertSuspend Term
t2)
  TTydef Var
x Poly Type
pty Maybe TydefInfo
mtd Term
t1 -> Var -> Poly Type -> Maybe TydefInfo -> Term -> Term
TTydef Var
x Poly Type
pty Maybe TydefInfo
mtd (Term -> Term
insertSuspend Term
t1)
  TBind Maybe Var
mx Maybe (Poly Type)
mty Maybe Requirements
mreq Term
c1 Term
c2 -> Maybe Var
-> Maybe (Poly Type) -> Maybe Requirements -> Term -> Term -> Term
TBind Maybe Var
mx Maybe (Poly Type)
mty Maybe Requirements
mreq Term
c1 (Term -> Term
insertSuspend Term
c2)
  TAnnotate Term
t1 Poly Type
ty -> Term -> Poly Type -> Term
TAnnotate (Term -> Term
insertSuspend Term
t1) Poly Type
ty
  -- Replace return or noop with suspend
  TApp (TConst Const
Return) Term
t1 -> Term -> Term
TSuspend Term
t1
  TConst Const
Noop -> Term -> Term
TSuspend Term
forall ty. Term' ty
TUnit
  -- Anything else: p => (__res__ <- p; suspend __res__)
  --
  -- Note that since we don't put type + reqs annotations on
  -- __res__, it won't get added to the type environment and hence
  -- won't be in scope for use after the suspend.  But we pick a
  -- weird unlikely-to-be-used name just to be safe.
  Term
t' -> Maybe Var
-> Maybe (Poly Type) -> Maybe Requirements -> Term -> Term -> Term
TBind (Var -> Maybe Var
forall a. a -> Maybe a
Just Var
"__res__") Maybe (Poly Type)
forall a. Maybe a
Nothing Maybe Requirements
forall a. Maybe a
Nothing Term
t' (Term -> Term
TSuspend (Var -> Term
forall ty. Var -> Term' ty
TVar Var
"__res__"))
 where
  thenSuspend :: Term
thenSuspend = Maybe Var
-> Maybe (Poly Type) -> Maybe Requirements -> Term -> Term -> Term
TBind Maybe Var
forall a. Maybe a
Nothing Maybe (Poly Type)
forall a. Maybe a
Nothing Maybe Requirements
forall a. Maybe a
Nothing Term
t (Term -> Term
TSuspend Term
forall ty. Term' ty
TUnit)