{-# LANGUAGE OverloadedStrings #-}

-- |
-- Module      :  Swarm.Language.Elaborate
-- Copyright   :  Brent Yorgey
-- Maintainer  :  byorgey@gmail.com
--
-- SPDX-License-Identifier: BSD-3-Clause
--
-- Term elaboration which happens after type checking.
module Swarm.Language.Elaborate where

import Control.Lens (transform, (%~), (^.))
import Swarm.Language.Syntax
import Swarm.Language.Types

-- | 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.  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 :: Syntax' Polytype -> Syntax' Polytype
elaborate :: Syntax' Polytype -> Syntax' Polytype
elaborate =
  -- Wrap all *free* variables in 'Force'.  Free variables must be
  -- referring to a previous definition, which are all wrapped in
  -- 'TDelay'.
  (forall ty. Traversal' (Syntax' ty) (Syntax' ty)
freeVarsS forall s t a b. ASetter s t a b -> (a -> b) -> s -> t
%~ \Syntax' Polytype
s -> forall ty. SrcLoc -> Term' ty -> ty -> Syntax' ty
Syntax' (Syntax' Polytype
s forall s a. s -> Getting a s a -> a
^. forall ty. Lens' (Syntax' ty) SrcLoc
sLoc) (forall ty. Syntax' ty -> Syntax' ty -> Term' ty
SApp Syntax' Polytype
sForce Syntax' Polytype
s) (Syntax' Polytype
s forall s a. s -> Getting a s a -> a
^. forall ty. Lens' (Syntax' ty) ty
sType))
    -- Now do additional rewriting on all subterms.
    forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a. Plated a => (a -> a) -> a -> a
transform Syntax' Polytype -> Syntax' Polytype
rewrite
 where
  rewrite :: Syntax' Polytype -> Syntax' Polytype
  rewrite :: Syntax' Polytype -> Syntax' Polytype
rewrite (Syntax' SrcLoc
l Term' Polytype
t Polytype
ty) = forall ty. SrcLoc -> Term' ty -> ty -> Syntax' ty
Syntax' SrcLoc
l (Term' Polytype -> Term' Polytype
rewriteTerm Term' Polytype
t) Polytype
ty

  rewriteTerm :: Term' Polytype -> Term' Polytype
  rewriteTerm :: Term' Polytype -> Term' Polytype
rewriteTerm = \case
    -- For recursive let bindings, rewrite any occurrences of x to
    -- (force x).  When interpreting t1, we will put a binding (x |->
    -- delay t1) in the context.
    SLet Bool
True LocVar
x Maybe Polytype
ty Syntax' Polytype
t1 Syntax' Polytype
t2 -> forall ty.
Bool
-> LocVar -> Maybe Polytype -> Syntax' ty -> Syntax' ty -> Term' ty
SLet Bool
True LocVar
x Maybe Polytype
ty (Var -> Syntax' Polytype -> Syntax' Polytype
wrapForce (LocVar -> Var
lvVar LocVar
x) Syntax' Polytype
t1) (Var -> Syntax' Polytype -> Syntax' Polytype
wrapForce (LocVar -> Var
lvVar LocVar
x) Syntax' Polytype
t2)
    -- Rewrite any recursive occurrences of x inside t1 to (force x).
    -- When a TDef is encountered at runtime its body will immediately
    -- be wrapped in a VDelay. However, to make this work we also need
    -- to wrap all free variables in any term with 'force' --- since
    -- any such variables must in fact refer to things previously
    -- bound by 'def'.
    SDef Bool
True LocVar
x Maybe Polytype
ty Syntax' Polytype
t1 -> forall ty.
Bool -> LocVar -> Maybe Polytype -> Syntax' ty -> Term' ty
SDef Bool
True LocVar
x Maybe Polytype
ty (Var -> Syntax' Polytype -> Syntax' Polytype
wrapForce (LocVar -> Var
lvVar LocVar
x) Syntax' Polytype
t1)
    -- Rewrite @f $ x@ to @f x@.
    SApp (Syntax' SrcLoc
_ (SApp (Syntax' SrcLoc
_ (TConst Const
AppF) Polytype
_) Syntax' Polytype
l) Polytype
_) Syntax' Polytype
r -> forall ty. Syntax' ty -> Syntax' ty -> Term' ty
SApp Syntax' Polytype
l Syntax' Polytype
r
    -- Leave any other subterms alone.
    Term' Polytype
t -> Term' Polytype
t

wrapForce :: Var -> Syntax' Polytype -> Syntax' Polytype
wrapForce :: Var -> Syntax' Polytype -> Syntax' Polytype
wrapForce Var
x = forall ty.
Var -> (Syntax' ty -> Syntax' ty) -> Syntax' ty -> Syntax' ty
mapFreeS Var
x (\s :: Syntax' Polytype
s@(Syntax' SrcLoc
l Term' Polytype
_ Polytype
ty) -> forall ty. SrcLoc -> Term' ty -> ty -> Syntax' ty
Syntax' SrcLoc
l (forall ty. Syntax' ty -> Syntax' ty -> Term' ty
SApp Syntax' Polytype
sForce Syntax' Polytype
s) Polytype
ty)

-- Note, TyUnit is not the right type, but I don't want to bother

sForce :: Syntax' Polytype
sForce :: Syntax' Polytype
sForce = forall ty. SrcLoc -> Term' ty -> ty -> Syntax' ty
Syntax' SrcLoc
NoLoc (forall ty. Const -> Term' ty
TConst Const
Force) (forall t. [Var] -> t -> Poly t
Forall [Var
"a"] (Fix TypeF -> Fix TypeF
TyDelay (Var -> Fix TypeF
TyVar Var
"a") Fix TypeF -> Fix TypeF -> Fix TypeF
:->: Var -> Fix TypeF
TyVar Var
"a"))