content: {-# LANGUAGE PatternSynonyms #-}
  
  -- |
  -- SPDX-License-Identifier: BSD-3-Clause
  --
  -- Helper functions for working with @Terms@ and @Syntax@
  module Swarm.Language.Syntax.Util (
    mkOp,
    mkOp',
    unfoldApps,
    mkTuple,
    unTuple,
  
    -- * Erasure
    erase,
    eraseS,
  
    -- * Term traversal
    freeVarsS,
    freeVarsT,
    freeVarsV,
    mapFreeS,
    locVarToSyntax',
    asTree,
    measureAstSize,
  ) where
  
  import Control.Lens (Traversal', para, universe, (%~), (^.), pattern Empty)
  import Control.Monad (void)
  import Data.Data (Data)
  import Data.List.NonEmpty (NonEmpty)
  import Data.List.NonEmpty qualified as NonEmpty
  import Data.Set qualified as S
  import Data.Tree
  import Swarm.Language.Syntax.AST
  import Swarm.Language.Syntax.Constants
  import Swarm.Language.Syntax.Loc
  import Swarm.Language.Syntax.Pattern
  import Swarm.Language.Var (Var)
  
  -- Setup for doctests
  
  -- $setup
  -- >>> import Control.Lens ((^.))
  -- >>> import Swarm.Language.Syntax.Constants
  -- >>> import Swarm.Language.Syntax.Loc
  -- >>> import Swarm.Language.Syntax.Pattern
  -- >>> import Swarm.Language.Syntax.AST
  
  -- | Make an infix operation (e.g. @2 + 3@) a curried function
  --   application (e.g. @((+) 2) 3@).
  mkOp :: Const -> (SrcLoc, t) -> Syntax -> Syntax -> Syntax
  mkOp c (opLoc, _) s1@(Syntax l1 _) s2@(Syntax l2 _) = Syntax newLoc newTerm
   where
    -- The new syntax spans all terms
    newLoc = l1 <> opLoc <> l2
    sop = Syntax opLoc (TConst c)
    newTerm = SApp (Syntax (l1 <> opLoc) $ SApp sop s1) s2
  
  -- | Make an infix operation, discarding any location information
  mkOp' :: Const -> Term -> Term -> Term
  mkOp' c t1 = TApp (TApp (TConst c) t1)
  
  -- | Turn function application chain into a list.
  --
  -- >>> syntaxWrap f = fmap (^. sTerm) . f . Syntax NoLoc
  -- >>> syntaxWrap unfoldApps (mkOp' Mul (TInt 1) (TInt 2)) -- 1 * 2
  -- TConst Mul :| [TInt 1,TInt 2]
  unfoldApps :: Syntax' ty -> NonEmpty (Syntax' ty)
  unfoldApps trm = NonEmpty.reverse . flip NonEmpty.unfoldr trm $ \case
    Syntax' _ (SApp s1 s2) _ _ -> (s2, Just s1)
    s -> (s, Nothing)
  
  -- | Create an appropriate `Term` out of a list of syntax nodes which
  --   were enclosed with parentheses (and separated by commas).
  mkTuple :: [Syntax] -> Term
  -- () = TUnit
  mkTuple [] = TUnit
  -- (x) = x, but record the fact that it was explicitly parenthesized,
  -- for better source location tracking
  mkTuple [x] = SParens x
  -- (x,y) = SPair
  mkTuple [x, y] = SPair x y
  -- (x,y,...) = recursively nested pairs.  Note that we do not assign
  -- source spans to the nested tuples since they don't really come from
  -- a specific place in the source.
  mkTuple (x : r) = SPair x (Syntax NoLoc (mkTuple r))
  
  -- | Decompose a nested tuple into a list of components.
  unTuple :: Syntax' ty -> [Syntax' ty]
  unTuple = \case
    Syntax' _ (SPair s1 s2) _ _ -> s1 : unTuple s2
    s -> [s]
  
  ------------------------------------------------------------
  -- Type erasure
  ------------------------------------------------------------
  
  -- | Erase the type annotations from a 'Syntax' or 'Term' tree.
  erase :: Functor t => t ty -> t ()
  erase = void
  
  -- | Erase all annotations from a 'Syntax' node, turning it into a
  --   bare 'Term'.
  eraseS :: Syntax' ty -> Term
  eraseS (Syntax' _ t _ _) = erase t
  
  ------------------------------------------------------------
  -- Free variable traversals
  ------------------------------------------------------------
  
  -- | Traversal over those subterms of a term which represent free
  --   variables.  The S suffix indicates that it is a `Traversal' over
  --   the `Syntax` nodes (which contain type and source location info)
  --   containing free variables inside a larger `Syntax` value.  Note
  --   that if you want to get the list of all `Syntax` nodes
  --   representing free variables, you can do so via @'toListOf'
  --   'freeVarsS'@.
  freeVarsS :: forall ty. Traversal' (Syntax' ty) (Syntax' ty)
  freeVarsS f = go S.empty
   where
    -- go :: Applicative f => Set Var -> Syntax' ty -> f (Syntax' ty)
    go bound s@(Syntax' l t ty cmts) = case t of
      TUnit -> pure s
      TConst {} -> pure s
      TDir {} -> pure s
      TInt {} -> pure s
      TAntiInt {} -> pure s
      TText {} -> pure s
      TAntiText {} -> pure s
      TAntiSyn {} -> pure s
      TBool {} -> pure s
      TRobot {} -> pure s
      TRef {} -> pure s
      TRequire {} -> pure s
      TStock {} -> pure s
      SRequirements x s1 -> rewrap $ SRequirements x <$> go bound s1
      TVar x
        | x `S.member` bound -> pure s
        | otherwise -> f s
      SLam x xty s1 -> rewrap $ SLam x xty <$> go (S.insert (lvVar x) bound) s1
      SApp s1 s2 -> rewrap $ SApp <$> go bound s1 <*> go bound s2
      SLet ls r x xty xpty xreq s1 s2 ->
        let bound' = S.insert (lvVar x) bound
         in rewrap $ SLet ls r x xty xpty xreq <$> go bound' s1 <*> go bound' s2
      STydef x xdef tdInfo t1 -> rewrap $ STydef x xdef tdInfo <$> go bound t1
      SPair s1 s2 -> rewrap $ SPair <$> go bound s1 <*> go bound s2
      SBind mx mty mpty mreq s1 s2 -> rewrap $ SBind mx mty mpty mreq <$> go bound s1 <*> go (maybe id (S.insert . lvVar) mx bound) s2
      SDelay s1 -> rewrap $ SDelay <$> go bound s1
      SRcd m -> rewrap $ SRcd <$> (traverse . traverse . traverse) (go bound) m
      SProj s1 x -> rewrap $ SProj <$> go bound s1 <*> pure x
      SAnnotate s1 pty -> rewrap $ SAnnotate <$> go bound s1 <*> pure pty
      SSuspend s1 -> rewrap $ SSuspend <$> go bound s1
      SParens s1 -> rewrap $ SParens <$> go bound s1
      TType {} -> pure s
     where
      rewrap s' = Syntax' l <$> s' <*> pure ty <*> pure cmts
  
  -- | Like 'freeVarsS', but traverse over the 'Term's containing free
  --   variables.  More direct if you don't need to know the types or
  --   source locations of the variables.  Note that if you want to get
  --   the list of all `Term`s representing free variables, you can do
  --   so via @'toListOf' 'freeVarsT'@.
  freeVarsT :: forall ty. Traversal' (Syntax' ty) (Term' ty)
  freeVarsT = freeVarsS . sTerm
  
  -- | Traversal over the free variables of a term.  Like 'freeVarsS'
  --   and 'freeVarsT', but traverse over the variable names themselves.
  --   Note that if you want to get the set of all free variable names,
  --   you can do so via @'Data.Set.Lens.setOf' 'freeVarsV'@.
  freeVarsV :: Traversal' (Syntax' ty) Var
  freeVarsV = freeVarsT . (\f -> \case TVar x -> TVar <$> f x; t -> pure t)
  
  -- | Apply a function to all free occurrences of a particular
  --   variable.
  mapFreeS :: Var -> (Syntax' ty -> Syntax' ty) -> Syntax' ty -> Syntax' ty
  mapFreeS x f = freeVarsS %~ (\t -> case t ^. sTerm of TVar y | y == x -> f t; _ -> t)
  
  -- | Transform the AST into a Tree datatype.  Useful for
  --   pretty-printing (e.g. via "Data.Tree.drawTree").
  asTree :: Data a => Syntax' a -> Tree (Syntax' a)
  asTree = para Node
  
  -- | Each constructor is a assigned a value of 1, plus
  --   any recursive syntax it entails.
  measureAstSize :: Data a => Syntax' a -> Int
  measureAstSize = length . filter (not . isNoop) . universe
  
  -- | Don't count "noop" nodes towards the code size.  They are usually
  --   inserted automatically, either in @{}@ or after a bare @def@.
  isNoop :: Syntax' a -> Bool
  isNoop = \case
    Syntax' _ (TConst Noop) _ _ -> True
    _ -> False
  
  locVarToSyntax' :: LocVar -> ty -> Syntax' ty
  locVarToSyntax' (LV s v) = Syntax' s (TVar v) Empty
  