{-------------------------------------------------------------------------------

        Copyright:              Bernie Pope 2004

        Module:                 Depend

        Description:            Compute the dependencies between functions
                                declared in a module. This information is
                                needed by the type checker.

        Primary Authors:        Bernie Pope

-------------------------------------------------------------------------------}

{-
    This file is part of baskell.

    baskell is free software; you can redistribute it and/or modify
    it under the terms of the GNU General Public License as published by
    the Free Software Foundation; either version 2 of the License, or
    (at your option) any later version.

    baskell is distributed in the hope that it will be useful,
    but WITHOUT ANY WARRANTY; without even the implied warranty of
    MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
    GNU General Public License for more details.

    You should have received a copy of the GNU General Public License
    along with baskell; if not, write to the Free Software
    Foundation, Inc., 59 Temple Place, Suite 330, Boston, MA  02111-1307  USA
-}

module Depend
   ( depend
   , printDepends
   )
   where

import AST
   ( Ident
   , Exp (..)
   , Decl (..)
   , Program (..)
   )

import qualified Data.Map as Map
   ( Map
   , empty
   , lookup
   , insert
   )

import Data.Graph
  ( scc
  , Vertex
  , Edge
  , Graph
  , buildG
  )

import Data.Tree
  ( flatten )

import Data.Maybe
   ( catMaybes )

--------------------------------------------------------------------------------

-- given a list of declarations compute the list of strongly connected
-- components.
depend :: [Decl] -> [[Decl]]
depend decls
   = convertToDecls $ map flatten $ scc $ dependGraph decls
   where
   convertToDecls :: [[Vertex]] -> [[Decl]]
   convertToDecls vertices
      = map (catMaybes . map (flip Map.lookup vertexMap)) vertices
   (vertexMap, _) = declMaps decls

-- mappings from vertex numbers to decls and vice versa
declMaps :: [Decl] -> (Map.Map Vertex Decl, Map.Map Ident Vertex)
declMaps decls
   = (intToDeclMap 0 decls, declToIntMap 0 decls)
   where
   intToDeclMap :: Vertex -> [Decl] -> Map.Map Vertex Decl
   intToDeclMap count [] = Map.empty
   intToDeclMap count (d:ds)
      = Map.insert count d (intToDeclMap (count + 1) ds)
   declToIntMap :: Vertex -> [Decl] -> Map.Map Ident Vertex
   declToIntMap count [] = Map.empty
   declToIntMap count (Sig {} : ds) =
      declToIntMap count ds
   declToIntMap count (Decl ident _body : ds)
      = Map.insert ident count (declToIntMap (count + 1) ds)

-- edges are immediate dependencies of the dependency graph
-- the dependency graph is computed from this
dependencies :: [Decl] -> [Edge]
dependencies decls
   = getDependencies identMap decls
   where
   (_, identMap) = declMaps decls
   getDependencies :: Map.Map Ident Int -> [Decl] -> [Edge]
   getDependencies identMap [] = []
   getDependencies identMap (Sig {} : decls)
      = getDependencies identMap decls
   getDependencies identMap (Decl ident body : decls)
      = makeEdges identMap ident body ++ getDependencies identMap decls
   makeEdges :: Map.Map Ident Int -> Ident -> Exp -> [Edge]
   makeEdges identMap ident body
      = case Map.lookup ident identMap of
           Nothing -> []
           Just index
              -> [ (index, d) | d <- dependents]
      where
      dependents = catMaybes [ Map.lookup i identMap | i <- freeVarsInExp body ]

-- compute the dependency graph of a list of decls
dependGraph :: [Decl] -> Graph
dependGraph decls
   = buildG bounds $ dependencies decls
   where
   bounds = (0, length decls - 1)

type IdentMap = Map.Map Ident ()

-- find the free variables in an expression
freeVarsInExp :: Exp -> [Ident]
freeVarsInExp exp = freeVars Map.empty Map.empty exp

freeVars :: IdentMap -> IdentMap -> Exp -> [Ident]
freeVars bound free (Var ident)
   = case Map.lookup ident bound of
        Nothing -> case Map.lookup ident free of
                      Nothing -> [ident]
                      Just () -> []
        Just () -> []
freeVars bound free (Lam ident body)
   = freeVars (Map.insert ident () bound) free body
freeVars bound free (LamStrict ident body)
   = freeVars (Map.insert ident () bound) free body
freeVars bound free (App e1 e2)
   = freeVars bound free e1 ++ freeVars bound free e2
freeVars bound free (Literal _) = []
freeVars bound free (Tuple exps)
   = concatMap (freeVars bound free) exps
freeVars bound free (Prim _name _impl) = []

-- find and print the strongly connected components of a program
printDepends :: Program -> IO ()
printDepends (Program decls)
   = print $ map (map declName) $ depend decls
   where
   declName :: Decl -> String
   declName (Sig i _)  = i
   declName (Decl i _) = i