-- |
-- Module      : Minigent.Environment
-- Copyright   : (c) Data61 2018-2019
--                   Commonwealth Science and Research Organisation (CSIRO)
--                   ABN 41 687 119 230
-- License     : BSD3
--
-- Definitions for global environments and local type contexts.
--
-- May be used qualified or unqualified.
{-# LANGUAGE GeneralizedNewtypeDeriving, DeriveFunctor #-}
module Minigent.Environment
  ( -- * Global Environments
    GlobalEnvironments (..)
  , emptyGlobalEnvironments
  , -- * Contexts
    Context
  , push
  , pop
  , use
  , topUsed
  , unused
  , factor
  , alter
  , reconcile
  ) where

import Minigent.Syntax

import Data.Maybe (mapMaybe)
import qualified Data.Map as M
import qualified Data.List as L

-- | Global environments are for top-level functions, both definitions and types.
data GlobalEnvironments
   = GlobalEnvs
   { defns :: M.Map FunName (VarName, Expr)
   , types :: M.Map FunName PolyType
   } deriving (Show)

emptyGlobalEnvironments = GlobalEnvs M.empty M.empty

-- | A value of type @Context a@ maps variable names to values of type @a@,
--   keeping track of the amount of times the variable is "used". It can be
--   thought of as a /stack/ of bindings, with below 'push' and 'pop' functions.
newtype Context a = Ctx [(VarName, Int, a)] deriving (Semigroup, Monoid, Functor, Show)

-- | Add a new binding for the given variable to the context, shadowing existing bindings.
push :: (VarName, a) -> Context a -> Context a
push (v, x) (Ctx m) = Ctx ( (v, 0,x) : m)

-- | Returns true if the topmost bound variable in the context has been "used".
topUsed :: Context a -> Bool
topUsed (Ctx ((_,u,_):_)) = u > 0
topUsed _ = error "topUsed called on empty context"

-- | Removes the topmost binding from the context. If the context is empty, does nothing.
pop :: Context a -> Context a
pop (Ctx xs) = Ctx (drop 1 xs)

-- | Given a variable name, look it up in the context, returning its associated value.
--   Also returns a new context where the variable's usage count has been increased,
--   and a boolean that is true if the variable had already been used before.
use :: (Show a) => VarName -> Context a -> (a, Bool, Context a)
use v (Ctx ls) = let (a, r, ls') = go ls in (a, r, Ctx ls')
  where
    go [] = error "use called on an out of scope variable"
    go (x@(var, usage, a) : xs)
      | v == var = (a, usage > 0, (var, usage + 1, a) : xs)
      | v /= var = let (a', i, xs') = go xs
                    in (a', i, x:xs')

-- | Returns all elements in the context that have been left unused.
unused :: Context a -> [a]
unused (Ctx ls) = mapMaybe (\(v,u,a) -> if u == 0 then Just a else Nothing) ls

-- | Extract specific variables from a context. Given a list of variable names
--   and an input context, the output will contain only those variables in
--   the input context mentioned in the list. 
--
--   This is used in let! expressions to extract observer variables from the
--   rest of the context.
factor :: [VarName] -> Context a -> Context a
factor vs (Ctx xs) = let as = go vs xs in Ctx as
  where
    go vs [] = []
    go vs ((v,u,a):xs)
      | v `elem` vs = (v,u,a): go (L.delete v vs) xs
      | otherwise   = go vs xs

-- | Alter specific bindings using the provided function.
--   
--   This is used in let! expressions to modify observer variable bindings.
alter :: [VarName] -> ((a, Int) -> (a, Int)) -> Context a -> Context a 
alter vs f (Ctx xs) = Ctx (map go xs)
  where 
    go (v,u,a)
      | v `elem` vs = let (a',u') = f (a,u) in (v,u',a')
      | otherwise   = (v,u,a)

-- | Given two contexts that are assumed to have the same variables
--   and elements, but may
--   differ in their usage counts, return a combined context where
--   the usage count is the maximum of the inputs, along with any elements
--   where the usage counts differed.
--
--   This is used in situations like if and case expressions, where
--   control flow branches, and the same linear variables must be used in both
--   branches.
reconcile :: Context a -> Context a -> ([a], Context a)
reconcile (Ctx xs) (Ctx ys) = let (as, zs) = go xs ys in (as, Ctx zs)
  where
    go [] [] = ([], [])
    go ((v,u,a):xs) ((_,u',_):ys)
      = let (as, zs) = go xs ys
            z        = (v,max u u',a)
            as'      = if u /= u' then a:as else as
         in (as', z:zs)