{-# OPTIONS_GHC -Wall #-}
module Canonicalize.Sort (definitions) where

import Control.Monad.State
import qualified Data.Graph as Graph
import qualified Data.Map as Map
import qualified Data.Maybe as Maybe
import qualified Data.Set as Set

import AST.Expression.General (Expr'(..))
import qualified AST.Expression.Canonical as Canonical
import qualified AST.Pattern as P
import qualified AST.Variable as V
import qualified Reporting.Annotation as A



-- STARTING POINT


definitions :: Canonical.Expr -> Canonical.Expr
definitions expression =
  evalState (reorder expression) Set.empty


free :: String -> State (Set.Set String) ()
free x =
  modify (Set.insert x)


freeIfLocal :: V.Canonical -> State (Set.Set String) ()
freeIfLocal (V.Canonical home name) =
  case home of
    V.Local -> free name
    V.TopLevel _ -> free name
    V.BuiltIn -> return ()
    V.Module _ -> return ()



-- REORDER EXPRESSIONS


reorder :: Canonical.Expr -> State (Set.Set String) Canonical.Expr
reorder (A.A ann expression) =
  A.A ann <$>
  case expression of
    -- Be careful adding and restricting freeVars
    Var var ->
      do  freeIfLocal var
          return expression

    Lambda pattern body ->
      uncurry Lambda <$> bindingReorder (pattern,body)

    Binop op leftExpr rightExpr ->
      do  freeIfLocal op
          Binop op <$> reorder leftExpr <*> reorder rightExpr

    Case expr cases ->
      Case <$> reorder expr <*> mapM bindingReorder cases

    Data name exprs ->
      do  free name
          Data name <$> mapM reorder exprs

    -- Just pipe the reorder though
    Literal _ ->
      return expression

    ExplicitList es ->
      ExplicitList <$> mapM reorder es

    App func arg ->
      App <$> reorder func <*> reorder arg

    If branches finally ->
      If
        <$> mapM (\(cond,branch) -> (,) <$> reorder cond <*> reorder branch) branches
        <*> reorder finally

    Access record field ->
      Access
        <$> reorder record
        <*> return field

    Update record fields ->
      Update
        <$> reorder record
        <*> mapM (\(field,expr) -> (,) field <$> reorder expr) fields

    Record fields ->
      Record
        <$> mapM (\(field,expr) -> (,) field <$> reorder expr) fields

    Cmd _ ->
      return expression

    Sub _ ->
      return expression

    OutgoingPort _ _ ->
      return expression

    IncomingPort _ _ ->
      return expression

    Program _ _ ->
      error "DANGER - Program AST nodes should not be in def sorting."

    SaveEnv _ _ ->
      return expression

    GLShader _ _ _ ->
      return expression

    -- Actually do some reordering
    Let defs body ->
      do  body' <- reorder body

          -- Sort defs into strongly connected components.This
          -- allows the programmer to write definitions in whatever
          -- order they please, we can still define things in order
          -- and generalize polymorphic functions when appropriate.
          sccs <- Graph.stronglyConnComp <$> buildDefGraph defs
          let defss = map Graph.flattenSCC sccs

          -- remove let-bound variables from the context
          forM_ defs $ \(Canonical.Def _ pattern _ _) -> do
              bound pattern
              mapM free (ctors pattern)

          let (A.A _ let') =
                foldr (\ds bod -> A.A ann (Let ds bod)) body' defss

          return let'


ctors :: P.Canonical -> [String]
ctors (A.A _ pattern) =
  case pattern of
    P.Var _ ->
        []

    P.Alias _ p ->
        ctors p

    P.Record _ ->
        []

    P.Anything ->
        []

    P.Literal _ ->
        []

    P.Data (V.Canonical home name) ps ->
        case home of
          V.Local -> name : rest
          V.TopLevel _ -> name : rest
          V.BuiltIn -> rest
          V.Module _ -> rest
        where
          rest = concatMap ctors ps


bound :: P.Canonical -> State (Set.Set String) ()
bound pattern =
  let
    boundVars =
      P.boundVarSet pattern
  in
    modify (\freeVars -> Set.difference freeVars boundVars)


bindingReorder
    :: (P.Canonical, Canonical.Expr)
    -> State (Set.Set String) (P.Canonical, Canonical.Expr)
bindingReorder (pattern, expr) =
  do  expr' <- reorder expr
      bound pattern
      mapM_ free (ctors pattern)
      return (pattern, expr')



-- BUILD DEPENDENCY GRAPH BETWEEN DEFINITIONS


-- This also reorders the all of the sub-expressions in the Def list.
buildDefGraph
    :: [Canonical.Def]
    -> State (Set.Set String) [(Canonical.Def, Int, [Int])]
buildDefGraph defs =
  do  pdefsDeps <- mapM reorderAndGetDependencies defs
      return $ realDeps (addKey pdefsDeps)
  where
    addKey :: [(Canonical.Def, [String])] -> [(Canonical.Def, Int, [String])]
    addKey =
      zipWith (\n (pdef,deps) -> (pdef,n,deps)) [0..]

    variableToKey :: (Canonical.Def, Int, [String]) -> [(String, Int)]
    variableToKey (Canonical.Def _ pattern _ _, key, _) =
      [ (var, key) | var <- P.boundVarList pattern ]

    variableToKeyMap :: [(Canonical.Def, Int, [String])] -> Map.Map String Int
    variableToKeyMap pdefsDeps =
      Map.fromList (concatMap variableToKey pdefsDeps)

    realDeps :: [(Canonical.Def, Int, [String])] -> [(Canonical.Def, Int, [Int])]
    realDeps pdefsDeps =
        map convert pdefsDeps
      where
        varDict =
          variableToKeyMap pdefsDeps

        convert (pdef, key, deps) =
            (pdef, key, Maybe.mapMaybe (flip Map.lookup varDict) deps)


reorderAndGetDependencies
    :: Canonical.Def
    -> State (Set.Set String) (Canonical.Def, [String])
reorderAndGetDependencies (Canonical.Def facts pattern expr mType) =
  do  globalFrees <- get
      -- work in a fresh environment
      put Set.empty
      expr' <- reorder expr
      localFrees <- get
      -- merge with global frees
      modify (Set.union globalFrees)
      return (Canonical.Def facts pattern expr' mType, Set.toList localFrees)