jstypes

{-# LANGUAGE ViewPatterns #-}
{-# LANGUAGE GeneralizedNewtypeDeriving #-}
import           Control.Applicative
import           Control.Monad.State
import           Data.List hiding (or)
import           Data.Map (Map)
import qualified Data.Map as M
import           Data.Maybe
import           Language.ECMAScript3
import           Language.ECMAScript3.Syntax.Annotations
import           Prelude hiding (or)

-- print the type of the given expression
analyze str =
  case parse parseStatement "<input>" str of
    Left e -> error (show e)
    Right e -> putStrLn $ str ++ " :: " ++ pretty (ty (evalState (runT (typifyStmt e)) M.empty))

-- simple data types describing js subset
data Type
  = Unknown
  | Int
  | Double
  | Bool
  | String
  | Null
  | Undefined
  | Or Type Type
  | Fun Type Type
  | Returns Type
  deriving (Show,Eq)

-- typifying monad
newtype T a = T { runT :: State (Map String Type) a }
  deriving (Monad,Functor,MonadState (Map String Type))

-- annotate an expression with its type(s)
typify :: Expression SourcePos -> T (Expression Type)
typify e =
  case e of
    IntLit _ i -> return (IntLit Int i)
    NumLit _ d
      | fromIntegral (round d) == d -> return (NumLit Int d)
      | otherwise -> return (NumLit Double d)
    InfixExpr _ op x y -> do
      x' <- typify x
      y' <- typify y
      return (InfixExpr (opType op (ty x') (ty y')) op x' y')
    BoolLit a b   -> return (BoolLit Bool b)
    StringLit _ b -> return (StringLit String b)
    NullLit _     -> return (NullLit Null)
    CondExpr _ cond x y -> do
      x' <- typify x
      y' <- typify y
      cond' <- typify cond
      return (CondExpr (or (ty x') (ty y')) cond' x' y')
    FuncExpr _ mname params stmts -> do
      stmts' <- mapM typifyStmt stmts
      return (FuncExpr (Fun Unknown (dropReturns (stmtsType (map ty stmts'))))
                       undefined
                       undefined
                       stmts')
    VarRef _ (Id _ i) -> do
      ty <- ref i
      let ty' = fromMaybe Undefined ty
      return $ VarRef ty' (Id ty' i)
    e -> error ("typify: " ++ show e)

-- drop any wrapping returns
dropReturns :: Type -> Type
dropReturns (Returns t) = t
dropReturns t = t

-- get the type of a list of statements
stmtsType :: [Type] -> Type
stmtsType = Returns . go . mapMaybe returns . shortCircuit where
  go (x:xs) = foldl or x xs
  go [] = Undefined

-- collect returns from nested or's and whatnot into top-level Return or nothing
returns (Returns x) = returns x <|> return x
returns (Or x y) = both <|> returns x <|> returns y
  where both = or <$> returns x <*> returns y
returns _ = Nothing

-- stop once a non-conditional return statement is found
shortCircuit = go where
  go (x@Returns{}:xs) = [x]
  go (x@(Or Returns{} Returns{}):xs) = [x]
  go (x:xs) = x : go xs
  go [] = []

-- typify a statement
typifyStmt :: Statement SourcePos -> T (Statement Type)
typifyStmt e =
  case e of
    ReturnStmt _ Nothing -> return (ReturnStmt (Returns Undefined) Nothing)
    ReturnStmt _ (Just x) -> do x' <- typify x
                                return (ReturnStmt (Returns (ty x')) (Just x'))
    ExprStmt _ e -> do x' <- typify e; return (ExprStmt (ty x') x')
    VarDeclStmt _ decls -> do decls' <- mapM typifyVarDecl decls
                              return $ VarDeclStmt Undefined decls'
    BlockStmt _ [] -> return (EmptyStmt Undefined)
    BlockStmt _ [stmt] -> typifyStmt stmt
    BlockStmt _ stmts -> do
      stmts' <- mapM typifyStmt stmts
      return (BlockStmt (stmtsType (map ty stmts')) stmts')
    IfSingleStmt sp pred pthen -> typifyStmt (IfStmt sp pred pthen (EmptyStmt sp))
    EmptyStmt _ -> return (EmptyStmt Undefined)
    IfStmt _ pred pthen pelse -> do
      pred' <- typify pred
      pthen' <- typifyStmt pthen
      pelse' <- typifyStmt pelse
      return $ IfStmt (or (ty pthen') (ty pelse'))
                      pred'
                      pthen'
                      pelse'
    e -> error ("typifyStmt: " ++ show e)

-- typify a var x = … declaration
typifyVarDecl :: VarDecl SourcePos -> T (VarDecl Type)
typifyVarDecl (VarDecl _ (Id _ i) exp) = do
  exp' <- maybe (return Nothing) (fmap Just . typify) exp
  let ty' = (maybe Undefined ty exp')
      ident' = Id ty' i
  bind i ty'
  return $ VarDecl Undefined ident' exp'

-- bind a variable with the given type
bind :: String -> Type -> T ()
bind i ty =
  modify $ M.insertWith (flip const) i ty

-- lookup the type of a variable reference
ref :: String -> T (Maybe Type)
ref i = do
  m <- get
  return $ M.lookup i m

-- get the type of an operator applied to two typed values
opType :: InfixOp -> Type -> Type -> Type
opType op x y =
  case op of
    -- crazy monoidal-but-casting-thing
    OpAdd        -> addOrConcat x y
    -- logical operations that result in boolean
    OpLT         -> Bool
    OpLEq        -> Bool
    OpGT         -> Bool
    OpGEq        -> Bool
    OpIn         -> Bool
    OpInstanceof -> Bool
    OpEq         -> Bool
    OpNEq        -> Bool
    OpStrictEq   -> Bool
    OpStrictNEq  -> Bool
    -- logical operations that result in the type of the last operand
    OpLAnd       -> or x y
    -- logical operations that result in… something else
    OpLOr        -> or x y
    -- hetero arithmetic operations
    OpMul        -> doi x y
    OpMod        -> doi x y
    OpSub        -> doi x y
    -- homo arithmetic operations
    OpDiv        -> Double
    -- not sure about these, will have to check the spec.
    OpLShift     -> Int
    OpSpRShift   -> Int
    OpZfRShift   -> Int
    OpBAnd       -> Int
    OpBXor       -> Int
    OpBOr        -> Int

-- double or integer
doi :: Type -> Type -> Type
doi Double _ = Double
doi _ Double = Double
doi Int _    = Int
doi _ Int    = Int
doi _ _      = Double

-- concatenation or addition, depends
addOrConcat :: Type -> Type -> Type
addOrConcat = go where
  -- let's call unknown + a = unknown
  go Unknown _ = Unknown
  go _ Unknown = Unknown
  -- string + a = string
  go _ String = String
  go String _ = String
  -- gonna go ahead and say that undefined combined with anything except string is undefined
  -- may re-think this later
  go Undefined _ = Undefined
  go _ Undefined = Undefined
  -- int->double promotion
  go Double Int = Double
  go Int Double = Double
  -- bools behave like numbers
  go Bool Int    = Int
  go Bool Double = Double
  go Int Bool    = Int
  go Double Bool = Double
  -- nulls also behave like numbers
  go Null Int    = Int
  go Null Double = Double
  go Int Null    = Int
  go Double Null = Double
  -- equalities
  go Int Int       = Int
  go Double Double = Double
  go Bool Bool     = Int
  go Null Null     = Int
  -- null + bool
  go Null Bool = Int
  go Bool Null = Int
  -- ors
  go (Or x y) Int      = omap (addOrConcat Int) x y
  go (Or x y) Double   = omap (addOrConcat Double) x y
  go (Or x y) Bool     = omap (addOrConcat Int) x y
  go (Or x y) Null     = omap (addOrConcat Int) x y
  -- ors flipped
  go Int (Or x y)      = omap (addOrConcat Int) x y
  go Double (Or x y)   = omap (addOrConcat Double) x y
  go Bool (Or x y)     = omap (addOrConcat Int) x y
  go Null (Or x y)     = omap (addOrConcat Int) x y
  -- ors general
  go (Or x y) (Or a b) = or (omap (addOrConcat (or a b)) x y)
                            (omap (addOrConcat (or x y)) a b)

-- apply a type transformation to both types in an Or t1 t2
omap :: (Type -> Type) -> Type -> Type -> Type
omap f x y = or (f x) (f y)

-- make an Or type, but collapse types in common, e.g. (Or x y) x == Or x y
or :: Type -> Type -> Type
or x y = go x y where
  go x y | x == y = x
         | otherwise = foldr1 Or (nub (collectOrs (Or x y)))

-- collect a nested tree of Or x (Or y z) into [x,y,z]
collectOrs :: Type -> [Type]
collectOrs = go where
  go (Or x y) = go x ++ go y
  go a        = [a]

-- get the type of an expression
ty :: HasAnnotation obj => obj Type -> Type
ty = getAnnotation

-- pretty print a type to a string
pretty :: Type -> String
-- pretty a = show a
pretty Unknown   = "?"
pretty Int       = "Int"
pretty Double    = "Double"
pretty Bool      = "Bool"
pretty String    = "String"
pretty Null      = "Null"
pretty Undefined = "Undefined"
pretty (Fun a b) = fparens (pretty a) ++ " -> " ++ fparens (pretty b)
pretty o@Or{} =
  "Or " ++ intercalate " " (map (parens . pretty) (collectOrs o)) where
pretty (Returns ty) = "Returns " ++ parens (pretty ty)

-- if parens are needed to disambiguate, add them
parens x | any (==' ') x = "(" ++ x ++ ")"
         | otherwise     = x

-- if parens are needed to disambiguate, add them
fparens x | isInfixOf " -> " x = "(" ++ x ++ ")"
          | otherwise     = x