jstypes.hs

-- Examples:

-- λ> analyze "1"
-- 1 :: Int
-- λ> analyze "(function(){})"
-- (function(){}) :: ? -> Undefined
-- λ> analyze "(function(){ var x = 1; })"
-- (function(){ var x = 1; }) :: ? -> Undefined
-- λ> analyze "(function(){ var x = 1; return x })"
-- (function(){ var x = 1; return x}) :: ? -> Int
-- λ> analyze "(function(){ var x = {}; x.y = 'hello!'; return x })"
-- (function(){ var x = {}; x.y = 'hello!'; return x }) :: ? -> {y :: String}
-- λ> analyze "(function(){ var x = {}; x.y = 'hello!'; return x || 123 })"
-- (function(){ var x = {}; x.y = 'hello!'; return x || 123 }) :: ? -> Or ({y :: String}) Int

{-# LANGUAGE ViewPatterns #-}
{-# LANGUAGE GeneralizedNewtypeDeriving #-}
import           Control.Applicative
import           Control.Arrow
import           Data.IORef
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)

main = catch (getLine >>= analyze >> main)
             (const (return ()))

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

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

-- typifying monad
newtype T a = T { runT :: StateT (Map String Type) IO a }
  deriving (Monad,Functor,MonadState (Map String Type),MonadIO,Alternative,Applicative)

-- 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')
    AssignExpr _ op lvalue exp ->
      assignExp op lvalue exp
    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)
    ObjectLit _ props -> do
      props <- mapM (\(prop,exp) -> do exp' <- typify exp
                                       let prop' = typifyProp prop (ty exp')
                                       return (prop',exp'))
                    props
      ref <- io $ newIORef (map (propName *** ty) props)
      return $ ObjectLit (Object ref) props
    DotRef _ exp (Id _ i) ->
      dotRef exp i
    e -> error ("typify: " ++ show e)

dotRef exp field = do
  exp' <- typify exp
  let ty' = ty exp'
  result <- lookupTypeField ty' field
  fieldtype <- case result of
    Nothing -> return Unknown -- TODO: and rebind the var
    Just fieldtype -> return fieldtype
  return $ DotRef fieldtype
                  exp'
                  (Id ty' field)

lookupTypeField ty field = do
  case ty of
    Object fields -> do
      fields' <- io $ readIORef fields
      case lookup field fields' of
        Just fieldtype -> return (Just fieldtype)
        _ -> return Nothing
    Or x y -> lookupTypeField x field <|> lookupTypeField y field
    _ -> return Nothing

propName p =
  case p of
    PropId _ (Id _ i) -> i
    PropString _ i -> i
    PropNum _ i -> show i

typifyProp p ty =
  case p of
    PropId _ (Id _ i) -> PropId ty (Id ty i)
    PropString _ i -> PropString ty i
    PropNum _ i -> PropNum ty i

assignExp OpAssign (LVar _ i) exp = do
  exp' <- typify exp
  let expty = ty exp'
  ty' <- do
    ty' <- ref i
    case ty' of
      Just t -> do when (t /= expty) $
--                   don't know which is better right now
                     rebind i (or t expty)
--                     rebind i expty
                   return t
      Nothing -> do bind i expty
                    return expty
  return $ AssignExpr expty OpAssign (LVar ty' i) exp'
assignExp OpAssign (LDot _ exp key) rhs = do
  exp' <- typify exp
  rhs' <- typify rhs
  updateType (ty exp') key (ty rhs')
  return $ AssignExpr (ty rhs') OpAssign (LDot (ty exp') exp' key) rhs'

updateType ty fieldname fieldtype =
  case ty of
    Object ref -> do
     io $ modifyIORef ref $ \pairs ->
       let sans = filter ((/= fieldname).fst) pairs
       in case lookup fieldname pairs of
            Nothing -> (fieldname,fieldtype) : pairs
            Just existingType -> (fieldname,unifyTypes existingType fieldtype) : sans
     return ()
    _ -> return () -- possibly warn here

unifyTypes :: Type -> Type -> Type
unifyTypes x y = or x y

-- 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

-- re-bind a variable with the given type
rebind :: String -> Type -> T ()
rebind i ty =
  modify $ M.insertWith (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 -> IO String
-- pretty a = show a
pretty Unknown   = return "?"
pretty Int       = return "Int"
pretty Double    = return "Double"
pretty Bool      = return "Bool"
pretty String    = return "String"
pretty Null      = return "Null"
pretty Undefined = return "Undefined"
pretty (Fun a b) = do
  pa <- (pretty a)
  pb <- (pretty b)
  return $ fparens pa ++ " -> " ++ fparens pb
pretty o@Or{} = do
  pps <- mapM pretty (collectOrs o)
  return $ "Or " ++ intercalate " " (map parens pps)
pretty (Returns ty) = do
  pp <- (pretty ty)
  return $ "Returns " ++ parens pp
pretty (Object pairs) = do
  pairs' <- readIORef pairs
  pps <- mapM keypair pairs'
  return $ "{" ++ intercalate "," pps ++ "}"
  where keypair (key,value) = do
          pp <- pretty value
          return $ key ++ " :: " ++ pp

-- 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

io = liftIO