{-
    Kaya - My favourite toy language.
    Copyright (C) 2004, 2005 Edwin Brady

    This file is distributed under the terms of the GNU General
    Public Licence. See COPYING for licence.
-}

module Language(module AbsSyntax, 
                module Context,
                module Language) where

----- Useful gadgets on the abstract syntax; also reexports syntax types.

import List
import Debug.Trace
import AbsSyntax
import Context
import Options

-- Merge somehow with 'Error' and tidy up name lookup functions.
data Lookup a = Got a
              | Ambiguous Name [Name]
              | UnknownName Name

instance Show a => Show (Lookup a) where
    show x = s' x
      where
        s' (Got a) = "Success: " ++ show a
        s' (Ambiguous n xs) = "Ambiguous name: " ++ showuser n ++
                              " (Could be " ++ showthings (sort xs) ++ ")"
        s' (UnknownName n) = "Unknown name: " ++ showuser n
	showthings [] = ""
	showthings [x] = showuser x
	showthings (x:xs) = showuser x ++ ", " ++ showthings xs


-- Lookup a name, bearing in mind namespaces.
-- Returns all possibilities, in the current module and others.
{-
lookupname :: Name -> -- current module
              Name -> -- name to lookup (possibly decorated)
	      [(Name,a)] -> [(Name,a)]
lookupname mod n gam = checkCurrent $ lu n (decorated n) gam [] where
    lu n _ [] acc = acc
    lu n True ((x,a):xs) acc | n==x = lu n True xs ((x,a):acc)
			     | otherwise = lu n True xs acc
    lu n False ((x,a):xs) acc | nameMatches n x = lu n False xs ((x,a):acc)
			      | otherwise = lu n False xs acc

    -- in theory, returns names in the current module if they exist, or
    -- all names if the name is not in the current module. But in the
    -- presence of ad-hoc overloading, I don't think this makes sense.
    checkCurrent xs = xs {- cc xs xs
    cc [] xs = xs
    cc ((NS m n,x):xs) _ | m == mod = (NS m n,x):(cc xs [])
    cc (_:ys) xs = cc ys xs -}

    decorated (NS _ _) = True
    decorated _ = False
    nameMatches n (NS _ a) = nameMatches n a
    nameMatches n x = n == x
-}


-- Lookup in the context (ignores namespace if no ambiguity, returns
-- fully qualified name)
-- FIXME: This should return a Lookup structure, to be more informative.
ctxtlookup :: Monad m => Name -> -- Current module
	                 Name -> Context -> 
                         Maybe Type -> -- type information to help disambiguate
                         Options -> -- compiler options
                         m (Name, Type)
ctxtlookup mod n ctx ty copts
    = do let opts = nub (lookupname mod n ctx)
	 let pub = filter accessible opts
	 let priv = opts \\ pub
         let pubs = (tymatch ty (nubnames pub))
         let depr = checkDepr pubs
	 returnName pubs priv depr

  where returnName _ _ True = case ty of 
                               (Just jty) -> fail $ "Can't use deprecated function " ++ showuser n ++ "(" ++ showargs jty ++ ")"
                               Nothing -> fail $ "Can't use deprecated function " ++ showuser n
        returnName [(x,(ty,opts))] _ False = return (x,ty)
        returnName [] [] False = fail $ "Unknown name " ++ showuser n
	returnName [] priv False = fail $ "Can't use private name: " ++ 
			         showuser n ++ " (" ++ showthings priv ++ ")"
	returnName pub _ False = fail $ "Ambiguous name: " ++ showuser n ++ 
			          " (Could be " ++ 
				  showthings pub ++ ")"
        
        checkDepr [] = False
-- only need to check first element
        checkDepr ((x,(ty,opts)):xs) = checkDepr' (elem DeprecatedFn opts)
-- something else will error first here
--        checkDepr _ = False

        checkDepr' False = False
        checkDepr' True = useDepr (elem DeprFail copts) (elem DeprWarn copts)

        useDepr False False = False
        useDepr True _ = True
-- FIXME: print a warning message below but succeed
        useDepr False True = False

	nubnames [] = []
	nubnames (f@(x,(ty,opts)):xs) | (x,ty) `elem` (map getpair xs) = nubnames xs
				      | otherwise = f:(nubnames xs)
        getpair (a,(b,c)) = (a,b)

        tymatch ty xs = case filter (matchesty ty) xs of
                           [] -> xs
                           x -> x

        matchesty (Just (Fn _ args _)) (x, ((Fn _ args2 _), opts))
            | (length args) <= (length args2) = and $ zipWith matchArgs args args2
            | otherwise = False
        matchesty _ _ = True

        matchArgs x y | x == y = True
        matchArgs (TyVar _) _ = True
        matchArgs _ (TyVar _) = True
        matchArgs (Array x) (Array y) = matchArgs x y
        matchArgs (TyApp t ts) (TyApp u us) = matchArgs t u &&
                                              (and $ zipWith matchArgs ts us)
        matchArgs (Fn _ args1 _) (Fn _ args2 _) 
            = (length args1 == length args2) && 
              and (zipWith matchArgs args1 args2)
        matchArgs _ _ = False

	-- A private name in the current module is accessible
      	accessible ((NS nmod _),_) | nmod == mod = True
        accessible (_,(_,fopts)) = elem Public fopts

        showthings xs = showStrs (sort $ map showsig xs)
        showsig (x,(ty,_)) = showuser x ++ "(" ++ showargs ty ++ ")"

        showargs (Fn _ args _) = showlist args
        showargs _ = ""

showStrs [] = ""
showStrs [x] = x
showStrs (x:xs) = x ++ ", " ++ showStrs xs

------------ Gadgets -------------

-- Return whether one type is "smaller" than another
-- X < Y if there is a (meaningful?) injection from X to Y.
-- There'll be a better way, if this table gets much bigger.
-- Should these be in a class PartialOrd?
tlt :: PrimType -> PrimType -> Bool
tlt Boolean Number = True
tlt Boolean StringType = True
tlt Character Number = True
tlt Character StringType = True
tlt Number RealNum = True
tlt Number StringType = True
tlt RealNum StringType = True
tlt _ _ = False

biggert :: Type -> Type -> Type
biggert (Prim x) (Prim y) | x `tlt` y = (Prim y)
biggert x y = x

mangling :: Type -> String
mangling t = "_" ++ mangling' t
mangling' (Fn _ args _) = "F" ++ concat (map mangling' args)
mangling' (Array arg) = "a" ++ mangling' arg
mangling' (User n) = show n
mangling' (TyApp n args) = mangling' n ++ concat (map mangling' args)
mangling' (TyVar _) = "P"
mangling' (Prim Number) = "i"
mangling' (Prim Character) = "c"
mangling' (Prim Boolean) = "b"
mangling' (Prim RealNum) = "f"
mangling' (Prim StringType) = "s"
mangling' (Prim Pointer) = "p"
mangling' (Prim Exception) = "e"
mangling' (Prim Void) = "v"
mangling' _ = ""

-- Get all of the type variables out of a type.

getTyVars :: Type -> [Name]
getTyVars (TyVar n) = [n]
getTyVars (TyApp f tys) = concat (map getTyVars (f:tys))
getTyVars (Array t) = getTyVars t
getTyVars (Fn _ tys t) = concat (map getTyVars tys) ++ getTyVars t
getTyVars _ = []

-- C Names need to be mangled with the type, for disambiguation of overloaded
-- functions

type Mangled = String

cname :: Name -> String -> Mangled
cname n mangle = show n ++ mangle

convert :: Type -> Type -> Bool
convert = (==)

checkConv :: Monad m => Type -> Type -> String -> m ()
checkConv x y err = if convert x y 
		     then return ()
		     else fail err

getType :: Monad m => Name -> [(Name,b)] -> m b
getType n ctxt = case (lookup n ctxt) of
		    Nothing -> fail $ "Unknown name gettype " ++ show n
		    (Just t) -> return t

getVars :: Type -> [Name]
getVars = nub.gv
    where gv (Fn ns ts t) = concat (map gv (t:ts))
	  gv (Array t) = gv t
	  gv (TyApp n ts) = concat (map gv (n:ts))
	  gv (TyVar n) = [n]
	  gv _ = []

numargs :: Type -> Int
numargs (Fn ns ts t) = length ts
numargs _ = 0


-- Give distinct type variables fresh names, so that independent variables
-- continue to be independent.

-- Okay, so this really ought to be called 'generalise', like in Algorithm W
-- for Hindley Milner inference. I couldn't remember the name at the time...

fudgevars :: Type -> Int -> (Type, Int)
fudgevars t next = let (vsmap,next') = newnames next (getUserVars t) in
		       (alpha vsmap t, next')
    where newnames n [] = ([],n)
	  newnames n (x:xs) = let (xsih,n') = newnames (n+1) xs in
				  ((x,MN ("TV",n)):xsih, n')
	  alpha vsmap (Prim x) = Prim x
	  alpha vsmap (Fn ns tys t) = Fn ns (map (alpha vsmap) tys)
				            (alpha vsmap t)
	  alpha vsmap (Array t) = Array (alpha vsmap t)
          alpha vsmap (User n) = User n
	  alpha vsmap (TyApp n tys) = tyapp (alpha vsmap n) 
                                            (map (alpha vsmap) tys)
--	  alpha vsmap (Syn n) = Syn n
	  alpha vsmap (TyVar x) = case lookup x vsmap of
				     Nothing -> TyVar x
				     (Just v) -> TyVar v
	  alpha vsmap UnknownType = UnknownType

getUserVars :: Type -> [Name]
getUserVars = nub.gv
    where gv (Fn ns ts t) = concat (map gv (t:ts))
	  gv (Array t) = gv t
	  gv (TyApp n ts) = concat (map gv (n:ts))
-- Actually it needs to be every variable so that constants get their
-- types inferred and generalised correctly - it was a faulty assumption
-- that all global names would have user defined type variables!
--	  gv (TyVar (UN n)) = [UN n]
	  gv (TyVar x) = [x]
	  gv _ = []

lvaltoexp :: RAssign -> Raw
lvaltoexp (RAName f l n) = RVar f l n
lvaltoexp (RAIndex f l lv r) = RIndex f l (lvaltoexp lv) r
lvaltoexp (RAField f l lv r) = RField f l (lvaltoexp lv) r

-- Lookup in the type context (ignores namespace if no ambiguity, returns
-- fully qualified type name)
typelookup :: Name -> -- Current module
	      Name -> Types -> Lookup (Name, TypeInfo)
typelookup mod t ti = returnName (nubnames (lookupname mod t ti))
   where returnName [x] = Got x
	 returnName [] = UnknownName t
--fail $ "Unknown type " ++ show t
	 returnName xs = Ambiguous t (map fst xs)
--fail $ "Ambiguous type name " ++ showuser t ++ 
--			        " (found " ++ showthings xs ++ ")"
	 showthings [] = ""
	 showthings [(x,_)] = showuser x
	 showthings ((x,_):xs) = showuser x ++ ", " ++ showthings xs

nubnames [] = []
nubnames (f@(x,_):xs) | x `elem` (map fst xs) = nubnames xs
	              | otherwise = f:(nubnames xs)

-- Lookup in the exception context (ignores namespace if no ambiguity, returns
-- fully qualified type name)
exceptlookup :: Name -> -- Current module
	        Name -> EContext -> Lookup (Name, [Type])
exceptlookup mod e ei = returnName (nubnames (lookupname mod e ei))
   where returnName [x] = Got x
	 returnName [] = UnknownName e
--fail $ "Unknown type " ++ show t
	 returnName xs = Ambiguous e (map fst xs)

-- Type normalisation; expand synonyms.

tyapp u [] = u
tyapp u ts = TyApp u ts

normalise :: Monad m => 
             Bool -> -- Unknown types should cause an error
             String -> Int -> Name -> Types -> Type -> m Type
normalise unk f l mod ti t = tn True [] t
 where
   tn top u (Fn ds ts t) 
       = do ts' <- mapM (tn top u) ts
	    t' <- tn top u t
	    return $ Fn ds ts' t'
   tn top u (Array t) = do t' <- tn top u t
		           return $ Array t'
   tn top u t@(User n) = 
       case typelookup mod n ti of
          un@(UnknownName _) -> 
               if unk 
                  then fail $ f ++ ":" ++ show l ++ ":" ++ show un
                  else return $ User (fixup mod n)
          am@(Ambiguous n xs) -> fail $ f ++ ":" ++ show l ++ ":" ++ show am
	  (Got (fqn, x)) -> applyTI top u fqn [] x
   tn top u (TyApp (User n) []) = tn top u (User n)
   tn top u t@(TyApp (User n) ts) = 
       case typelookup mod n ti of
          un@(UnknownName _) -> 
               do ts' <- mapM (tn False u) ts
                  return $ tyapp (User (fixup mod n)) ts'
          am@(Ambiguous n xs) -> fail $ f ++ ":" ++ show l ++ ":" ++ show am
	  (Got (fqn, x)) -> applyTI top u fqn ts x
   tn top u (TyApp n ts) =
       do ts' <- mapM (tn top u) ts
          n' <- tn top u n 
          return $ tyapp n' ts'
   tn top u rest = return rest

   fixup m fqn@(NS _ _) = fqn
   fixup m n = (NS m n)
--   fixup m n = n

   applyTI top u n ts (UserData as)
	| top && length ts < length as 
	    = fail $ f ++ ":" ++ show l ++ ":" ++
	       "Type " ++ showuser n ++ " has too few parameters"
	| top && length ts > length as 
	    = fail $ f ++ ":" ++ show l ++ ":" ++
	       "Type " ++ showuser n ++ " has too many parameters"
	| otherwise = do ts' <- mapM (tn False u) ts
			 return $ tyapp (User n) ts'
   -- Replace type with 't', replacing instances of as inside t with
   -- corresponding instances of ts.
   -- That probably makes no sense.
   applyTI top u n ts (Syn as t)
	| length ts < length as
	    = fail $ f ++ ":" ++ show l ++ ":" ++
	       "Type synonym " ++ showuser n ++ " has too few parameters"
	| length ts > length as
	    = fail $ f ++ ":" ++ show l ++ ":" ++
	       "Type synonym " ++ showuser n ++ " has too many parameters"
	| otherwise = if elem n u 
		       then fail $ f ++ ":" ++ show l ++ ":" ++
			     "Cycle in type synonyms " ++ 
			     showsyns u
		       else do st <- substty (zip as ts) t
			       tn top (n:u) st
     where showsyns [n] = showuser n
	   showsyns (n:ns) = showuser n ++ ", " ++ showsyns ns

   applyTI top u n ts Private = fail $ f ++ ":" ++ show l ++ 
			     "Can't use private type " ++ showuser n

   substty tmap (TyVar n) = case lookup n tmap of
			      Nothing -> fail $ "Shouldn't happen" ++ show tmap
			      (Just t) -> return t
   substty tmap (Array t) = do t' <- substty tmap t
			       return $ Array t'
   substty tmap (Fn ds as r) = do as' <- mapM (substty tmap) as
				  r' <- substty tmap r
				  return $ Fn ds as' r'
   substty tmap (TyApp n ts) = do ts' <- mapM (substty tmap) ts
                                  n' <- substty tmap n
				  return $ tyapp n' ts'
   substty _ rest = return rest

-- Fold constants in a raw term
-- TODO/FIXME: Check bounds?
cfold :: Raw -> Raw
cfold r@(RInfix f l op (RConst _ _ (Num x)) (RConst _ _ (Num y)))
    = case (foldint op x y) of
          Just c -> RConst f l c
	  Nothing -> r
cfold r@(RInfix f l op (RConst _ _ (Re x)) (RConst _ _ (Re y)))
    = case (foldreal op x y) of
          Just c -> RConst f l c
	  Nothing -> r
cfold r@(RUnary f l op (RConst _ _ (Num x)))
    = case (foldunint op x) of
          Just c -> RConst f l c
	  Nothing -> r
cfold r@(RUnary f l op (RConst _ _ (Re x)))
    = case (foldunreal op x) of
          Just c -> RConst f l c
	  Nothing -> r
cfold r = r

getConst (RConst _ _ c) = c

foldint :: Op -> Int -> Int -> Maybe Const
foldint Plus x y = Just $ Num (x+y)
foldint Minus x y = Just $ Num (x-y)
foldint Times x y = Just $ Num (x*y)
-- TODO: Should be compile time error
foldint Divide x 0 = Nothing
foldint Divide x y = Just $ Num (x `div` y)
foldint Modulo x y = Just $ Num (x `mod` y)
foldint Power x y = Just $ Num (floor ((fromIntegral x)**(fromIntegral y)))
foldint Equal x y = Just $ Bo (x==y)
foldint NEqual x y = Just $ Bo (x/=y)
foldint OpLT x y = Just $ Bo (x<y)
foldint OpGT x y = Just $ Bo (x>y)
foldint OpLE x y = Just $ Bo (x<=y)
foldint OpGE x y = Just $ Bo (x>=y)
foldint _ x y = Nothing 

foldreal :: Op -> Double -> Double -> Maybe Const
foldreal Plus x y = Just $ Re (x+y)
foldreal Minus x y = Just $ Re (x-y)
foldreal Times x y = Just $ Re (x*y)
-- TODO: Should be compile time error
foldreal Divide x 0 = Nothing
foldreal Divide x y = Just $ Re (x/y)
foldreal Power x y = Just $ Re (x**y)
foldreal Equal x y = Just $ Bo (x==y)
foldreal NEqual x y = Just $ Bo (x/=y)
foldreal OpLT x y = Just $ Bo (x<y)
foldreal OpGT x y = Just $ Bo (x>y)
foldreal OpLE x y = Just $ Bo (x<=y)
foldreal OpGE x y = Just $ Bo (x>=y)
foldreal _ x y = Nothing 

foldunint :: UnOp -> Int -> Maybe Const
foldunint Neg x = Just $ Num (-x)
foldunint _ _ = Nothing

foldunreal :: UnOp -> Double -> Maybe Const
foldunreal Neg x = Just $ Re (-x)
foldunreal _ _ = Nothing

-- compile time coercions (only do the ones which the type checker accepts)

mkInt :: Const -> Int
mkInt (Num x) = x
mkInt (Ch c) = fromEnum c
mkInt (Bo False) = 0
mkInt (Bo True) = 1
mkInt (Re x) = fromEnum x
mkInt (Str str) = read str
mkInt e = error $ "This can't happen: mkInt " ++ show e ++ " - please report"

mkFloat :: Const -> Double
mkFloat (Num x) = toEnum x
mkFloat (Ch c) = toEnum (fromEnum c)
mkFloat (Bo False) = 0.0
mkFloat (Bo True) = 1.0
mkFloat (Re x) = x
mkFloat (Str str) = read str
mkFloat e 
    = error $ "This can't happen: mkFloat " ++ show e ++ " - please report"

-- Apply a function (non-recursively) to every sub expression,
-- applying a different function to metavariables
-- (I don't know if this is actually that useful, but it is used by the 
-- optimiser...)
mapsubexpr f mf expr = app expr
  where app (Metavar fl l x) = mf fl l x
        app (Lambda ivs args e) = Lambda ivs args (f e)
	app (Closure args t e) = Closure args t (f e)
	app (Bind n ty e1 e2) = Bind n ty (f e1) (f e2)
	app (Declare fn l n t e) = Declare fn l n t (f e)
	app (Return e) = Return (f e)
	app (Assign a e) = Assign (aapply a) (f e)
	app (AssignOp op a e) = AssignOp op (aapply a) (f e)
	app (Seq a b) = Seq (f a) (f b)
	app (Apply fn as) = Apply (f fn) (applys as)
	app (ConApply fn as) = ConApply (f fn) (applys as)
	app (Partial b fn as i) = Partial b (f fn) (applys as) i
	app (Foreign ty n es) = Foreign ty n 
			        (zip (applys (map fst es)) (map snd es))
	app (While e b) = While (f e) (f b)
	app (DoWhile e b) = DoWhile (f e) (f b)
	app (For i nm j a e1 e2) = For i nm j (aapply a) (f e1) (f e2)
	app (TryCatch t e fl fin) = TryCatch (f t) (f e) (f fl) (f fin)
	app (Throw e) = Throw (f e)
	app (Except e i) = Except (f e) (f i)
	app (Infix op x y) = Infix op (f x) (f y)
	app (CmpStr op x y) = CmpStr op (f x) (f y)
	app (CmpExcept op x y) = CmpExcept op (f x) (f y)
	app (RealInfix op x y) = RealInfix op (f x) (f y)
	app (Append x y) = Append (f x) (f y)
        app (AppendChain xs) = AppendChain (applys xs)
	app (Unary op x) = Unary op (f x)
	app (RealUnary op x) = RealUnary op (f x)
	app (Coerce t1 t2 x) = Coerce t1 t2 (f x)
	app (Case e as) = Case (f e) (altapp as)
	app (If a t e) = If (f a) (f t) (f e)
	app (Index a b) = Index (f a) (f b)
	app (Field e n i j) = Field (f e) n i j
	app (ArrayInit as) = ArrayInit (applys as)
	app (Annotation a e) = Annotation a (f e)
        app (Length s) = Length (f s)
	app x = x

        aapply (AIndex a e) = AIndex (aapply a) (f e)
	aapply (AField a n i j) = AField (aapply a) n i j
	aapply x = x

        applys [] = []
	applys (x:xs) = (f x) : (applys xs)

        altapp [] = []
	altapp ((Alt i j es e):as) 
	    = (Alt i j (applys es) (f e)):(altapp as)
	altapp ((ArrayAlt es e):as) 
	    = (ArrayAlt (applys es) (f e)):(altapp as)
	altapp ((Default e):as) 
	    = (Default (f e)):(altapp as)
	altapp ((ConstAlt pt c e):as) 
	    = (ConstAlt pt c (f e)):(altapp as)

-- Fold a function across all sub expressions.
-- Applies 'f' to the subexpression, and uses 'com' to combine the
-- result across all sub expressions.
foldsubexpr :: (Expr n -> a) -> (a -> a -> a) -> a -> Expr n -> a
foldsubexpr f com def expr = app expr
  where app (Lambda ivs args e) = f e
	app (Closure args t e) = f e
	app (Bind n ty e1 e2) = (f e1) `com` (f e2)
	app (Declare fn l n t e) = (f e)
	app (Return e) = (f e)
	app (Assign a e) = (aapply a) `com` (f e)
	app (AssignOp op a e) = (aapply a) `com` (f e)
	app (Seq a b) = (f a) `com` (f b)
	app (Apply fn as) = (f fn) `com` (applys as)
	app (ConApply fn as) = (f fn) `com` (applys as)
	app (Partial b fn as i) = (f fn) `com` (applys as)
	app (Foreign ty n es) = applys (map fst es)
	app (While e b) = (f e) `com` (f b)
	app (DoWhile e b) = (f e) `com` (f b)
	app (For i nm j a e1 e2) = (aapply a) `com` (f e1) `com` (f e2)
	app (TryCatch t e fl fin) = (f t) `com` (f e) `com` 
				    (f fl) `com` (f fin)
	app (Throw e) = (f e)
	app (Except e i) = (f e) `com` (f i)
	app (Infix op x y) = (f x) `com` (f y)
	app (CmpStr op x y) = (f x) `com` (f y)
	app (CmpExcept op x y) = (f x) `com` (f y)
	app (RealInfix op x y) = (f x) `com` (f y)
	app (Append x y) = (f x) `com` (f y)
	app (Unary op x) = (f x)
	app (RealUnary op x) = (f x)
	app (Coerce t1 t2 x) = (f x)
	app (Case e as) = (f e) `com` (altapp as)
	app (If a t e) = (f a) `com` (f t) `com` (f e)
	app (Index a b) = (f a) `com` (f b)
	app (Field e n i j) = (f e)
	app (ArrayInit as) = (applys as)
	app (Annotation a e) = (f e)
	app x = def

        aapply (AIndex a e) = (aapply a) `com` (f e)
	aapply (AField a n i j) = (aapply a)
	aapply x = def

        applys [] = def
	applys (x:xs) = (f x) `com` (applys xs)

        altapp [] = def
	altapp ((Alt i j es e):as) 
	    = (applys es) `com` (f e) `com` (altapp as)
	altapp ((ArrayAlt es e):as) 
	    = (applys es) `com` (f e) `com` (altapp as)
	altapp ((Default e):as) 
	    = (f e) `com` (altapp as)
	altapp ((ConstAlt pt c e):as) 
	    = (f e) `com` (altapp as)

locsUsed :: Expr n -> [Int]
locsUsed (Loc i) = [i]
locsUsed (Lambda _ _ e) = locsUsed e
locsUsed (Closure _ _ e) = locsUsed e
locsUsed (Bind _ _ e1 e2) = locsUsed e1 ++ locsUsed e2
locsUsed (Declare _ _ _ _ e) = locsUsed e
locsUsed (Return e) = locsUsed e
locsUsed (Assign a e) = alocsUsed a ++ locsUsed e
locsUsed (AssignOp _ a e) = alocsUsed a ++ locsUsed e
locsUsed (AssignApp a e) = alocsUsed a ++ locsUsed e
locsUsed (Seq x y) = locsUsed x ++ locsUsed y
locsUsed (Apply f as) = locsUsed f ++ concat (map locsUsed as)
locsUsed (ConApply f as) = locsUsed f ++ concat (map locsUsed as)
locsUsed (Partial b f as _) = locsUsed f ++ concat (map locsUsed as)
locsUsed (Foreign _ _  as) = concat (map locsUsed (map fst as))
locsUsed (While x y) = locsUsed x ++ locsUsed y
locsUsed (DoWhile x y) = locsUsed x ++ locsUsed y
locsUsed (For _ _ _ a x y) = alocsUsed a ++ locsUsed x ++ locsUsed y
locsUsed (TryCatch x y z w) = locsUsed x ++ locsUsed y ++
                              locsUsed z ++ locsUsed w
locsUsed (NewTryCatch x c) = locsUsed x ++ concat (map clocsUsed c)
locsUsed (Throw x) = locsUsed x
locsUsed (Except x y) = locsUsed x ++ locsUsed y
locsUsed (NewExcept xs) = concat (map locsUsed xs)
locsUsed (Infix _ x y) = locsUsed x ++ locsUsed y
locsUsed (RealInfix _ x y) = locsUsed x ++ locsUsed y
locsUsed (CmpExcept _ x y) = locsUsed x ++ locsUsed y
locsUsed (CmpStr _ x y) = locsUsed x ++ locsUsed y
locsUsed (Append x y) = locsUsed x ++ locsUsed y
locsUsed (AppendChain xs) = concat (map locsUsed xs)
locsUsed (Unary _ x) = locsUsed x
locsUsed (RealUnary _ x) = locsUsed x
locsUsed (Coerce _ _ x) = locsUsed x
locsUsed (Case e cs) = locsUsed e ++ concat (map caseLocsUsed cs)
locsUsed (If t x y) = locsUsed t ++ locsUsed x ++ locsUsed y
locsUsed (Index x y) = locsUsed x ++ locsUsed y
locsUsed (Field e _ _ _) = locsUsed e
locsUsed (ArrayInit xs) = concat (map locsUsed xs)
locsUsed (Annotation _ e) = locsUsed e
locsUsed _ = []

alocsUsed (AIndex a e) = alocsUsed a ++ locsUsed e
alocsUsed (AField a _ _ _) = alocsUsed a
alocsUsed _ = []

clocsUsed (Catch (Right e) h) = locsUsed e ++ locsUsed h
clocsUsed (Catch (Left (n,es)) h) = concat (map locsUsed es) ++ locsUsed h

caseLocsUsed (Default e) = locsUsed e
caseLocsUsed (Alt _ _ es e) = concat (map locsUsed es) ++ locsUsed e
caseLocsUsed (ArrayAlt es e) = concat (map locsUsed es) ++ locsUsed e
caseLocsUsed (ConstAlt _ _ e) = locsUsed e

-- get a list of all variables which are modified (i.e. assigned to,
-- or passed to a function as a var argument) in a block

modified expr = modVar expr -- foldsubexpr modVar (++) [] expr

modVar :: Expr n -> [Int]
modVar (Lambda _ _ e) = modVar e
modVar (Closure _ _ e) = modVar e
modVar (Bind _ _ v e) = modVar v ++ modVar e
modVar (Declare _ _ _ _ e) = modVar e
modVar (Assign lval ex) = inLval lval ++ modVar ex
modVar (AssignOp _ lval ex) = inLval lval ++ modVar ex
modVar (AssignApp lval ex) = inLval lval ++ modVar ex
modVar (Seq x y) = modVar x ++ modVar y
modVar (Annotation a e) = modVar e
modVar (For _ _ _ l a b) = inLval l ++ modVar b ++ modVar a
modVar (While t e) = modVar t ++ modVar e
modVar (DoWhile e t) = modVar t ++ modVar e
modVar (NewExcept es) = concat (map modVar es)
modVar (Except e t) = modVar t ++ modVar e
modVar (Infix _ x y) = modVar x ++ modVar y
modVar (RealInfix _ x y) = modVar x ++ modVar y
modVar (Unary _ x) = modVar x
modVar (RealUnary _ x) = modVar x
modVar (Coerce _ _ e) = modVar e
modVar (Index x y) = modVar x ++ modVar y
modVar (Field x _ _ _) = modVar x
modVar (ArrayInit es) = concat (map modVar es)
modVar (NewTryCatch e cs) = modVar e ++ concat (map cmodVar cs)
-- This ought to check if the arg is a var arg.
modVar (Apply _ args) = concat (map locsUsed args)
modVar (ConApply _ args) = concat (map locsUsed args)
-- May be over conservative, but unlikely to happen much in practice.
modVar (Partial _ _ args _) = concat (map locsUsed args)
modVar (Foreign _ _ args) = concat (map locsUsed (map fst args))
modVar (Case e alts) = modVar e ++ concat (map inAlt alts)
modVar (If x t e) = modVar x ++ modVar t ++ modVar e
modVar (Return e) = modVar e
modVar (Throw e) = modVar e
modVar x = []

isArg (Loc i) = [i]
isArg _ = []

inAlt (Alt _ _ args e) = modVar e ++ concat (map isArg args)
inAlt (ArrayAlt args e) = modVar e ++ concat (map isArg args)
inAlt (ConstAlt _ _ e) = modVar e
inAlt (Default e) = modVar e

inLval (AName i) = [i]
inLval (AIndex lval _) = inLval lval
inLval (AField lval _ _ _) = inLval lval
inLval _ = []

cmodVar (Catch (Right nm) e) = modVar nm ++ modVar e
cmodVar (Catch (Left (_,es)) e) = concat (map modVar es) ++ modVar e

-- Functions which the compiler assumes to exist

eqfun = NS (UN "Builtins") (UN "equal")
eqmangle = mangling (Fn [] [TyVar (UN "a"), TyVar (UN "a")] (Prim Boolean))

sizefn = NS (UN "Builtins") (UN "size")
sizemangle = mangling (Fn [] [Array (TyVar (UN "a"))] (Prim Number))

missingCase = NS (UN "Builtins") (UN "Missing_Case")
missingCaseMangling = mangling (Fn [] [] (Prim Exception));

pmAssignFail = NS (UN "Builtins") (UN "Pattern_Matching_Assignment_Failure")
pmAssignFailMangling = mangling (Fn [] [] (Prim Exception));

exitfun = (NS (UN "Builtins") (UN "exit"))
exitmangle = mangling (Fn [] [(Prim Number)] (Prim Void))

pushfun = (NS (UN "Array") (UN "push"))

putstrlnfun = (NS (UN "Prelude") (UN "putStrLn"))
putstrlnmangle = mangling (Fn [] [(Prim StringType)] (Prim Void))

backtracefun = (NS (UN "Builtins") (UN "exceptionBacktrace"))
backtracemangle = mangling (Fn [] [(Prim Exception)] (Prim Void))

tappFun = (NS (UN "Multicore") (UN "tapply"))
tappMangle i = mangling (Fn [] (tappArgs 0 i) (TyVar (UN "a")))
   where tappArgs i num 
             | i == num = []
             | otherwise = (TyVar (UN ("a"++show i))):(tappArgs (i+1) num)

texecFun = (NS (UN "Multicore") (UN "texec"))
texecMangle i = mangling (Fn [] (tappArgs 0 i) (Prim Void))
   where tappArgs i num 
             | i == num = []
             | otherwise = (TyVar (UN ("a"++show i))):(tappArgs (i+1) num)

dumpFun = (NS (UN "Reflect") (UN "dump"))
dumpMangle = mangling (Fn [] [(TyVar (UN "a"))] (Prim Void))