content: module Kit.Compiler.Ir.ExprToIr (typedToIr, newIrContext) where
  
  import Control.Exception
  import Control.Monad
  import Data.Function
  import Data.Mutable
  import Data.List
  import Data.Maybe
  import Kit.Ast
  import Kit.Compiler.Context
  import Kit.Compiler.Ir.FindUnderlyingType
  import Kit.Compiler.Ir.PatternMatchToIr
  import Kit.Compiler.Ir.StringCompare
  import Kit.Compiler.Module
  import Kit.Compiler.TypeContext
  import Kit.Error
  import Kit.HashTable
  import Kit.Ir
  import Kit.NameMangling
  import Kit.Str
  
  data IrTempInfo = IrTempInfo {
    irTempType :: ConcreteType,
    irTempExpr :: Maybe TypedExpr,
    irTempOrder :: Int
  }
  
  data IrContext = IrContext {
    ictxTemps :: HashTable (TypedExpr, Span) Str,
    ictxTempInfo :: HashTable Int IrTempInfo,
    ictxNextTempId :: IORef Int,
    ictxChildNo :: Int,
    ictxIsLoopCondition :: Bool
  }
  
  newIrContext :: IO IrContext
  newIrContext = do
    d          <- h_new
    tempInfo   <- h_new
    nextTempId <- newRef 1
    return $ IrContext
      { ictxTemps           = d
      , ictxTempInfo        = tempInfo
      , ictxNextTempId      = nextTempId
      , ictxChildNo         = 0
      , ictxIsLoopCondition = False
      }
  
  makeIrTempVar :: IrContext -> IO Int
  makeIrTempVar ictx = do
    next <- readRef $ ictxNextTempId ictx
    modifyRef (ictxNextTempId ictx) ((+) 1)
    return next
  
  irTempName :: Int -> Str
  irTempName i = s_concat ["__tmp", s_pack $ show i]
  
  maybeTypedToIr ctx ictx mod e = case e of
    Just ex -> do
      t <- typedToIr ctx ictx mod ex
      return $ Just t
    Nothing -> return Nothing
  
  typedToIr :: CompileContext -> IrContext -> Module -> TypedExpr -> IO IrExpr
  typedToIr ctx ictx mod e@(TypedExpr { tExpr = et, tPos = pos, inferredType = t })
    = do
      let converter' = converter
            (typedToIr ctx ictx mod)
            (\pos -> findUnderlyingType ctx mod (Just pos))
      let paramConverter = \p -> converter'
      let r x = typedToIr ctx ictx mod x
      let maybeR x = case x of
            Just x -> do
              r' <- r x
              return $ Just r'
            Nothing -> return Nothing
      f <- findUnderlyingType ctx mod (Just pos) t
  
      case et of
        (Block children) -> do
          ictx     <- newIrContext
          children <- forMWithErrors (zip [0 ..] children)
            $ \(n, child) -> typedToIr ctx (ictx { ictxChildNo = n }) mod child
          lastId    <- readRef (ictxNextTempId ictx)
          tempDecls <- forM [1 .. lastId - 1] $ \i -> do
            temp@(IrTempInfo { irTempType = tempType, irTempExpr = tempDefault, irTempOrder = tempOrder }) <-
              h_get (ictxTempInfo ictx) i
            x <- typedToIr ctx ictx mod $ makeExprTyped
              (LocalVarDeclaration (Var ([], irTempName i))
                                   tempType
                                   True
                                   tempDefault
              )
              tempType
              pos
            return $ (tempOrder * 2, x)
          children <- return
            [ (n * 2 + 1, child) | (n, child) <- zip [0 ..] children ]
          allChildren <-
            return
            $ map snd
            $ sortBy (compare `on` fst)
            $ [ (n, xi)
              | (n, x) <- (tempDecls ++ children)
              , xi     <- case x of
                IrCompound xs -> xs
                _             -> [x]
              ]
          -- interleave temp decls with block children
          return $ IrBlock $ allChildren
        (Temp x) -> do
          existing <- h_lookup (ictxTemps ictx) (x, tPos x)
          case existing of
            Just s  -> return $ IrIdentifier ([], s)
            Nothing -> do
              let assignable = case f of
                    CArray _ _ -> False
                    _          -> True
              if not assignable
                then r x
                else do
                  r x
                  nextId <- makeIrTempVar ictx
                  h_insert (ictxTemps ictx) (x, tPos x) (irTempName nextId)
                  h_insert (ictxTempInfo ictx) nextId $ IrTempInfo
                    { irTempType  = t
                    , irTempExpr  = Just x
                    , irTempOrder = ictxChildNo ictx
                    }
                  return $ IrIdentifier ([], irTempName nextId)
        (Using   _            e1) -> r e1
        (Meta    m            e1) -> r e1
        (Literal (IntValue v) _ ) -> do
          return $ IrCast (IrLiteral (IntValue v) f) f
        (Literal l@(FloatValue v) _) -> case f of
          -- FIXME: this is better handled at the code generation stage now that we have literal types
          BasicTypeFloat 32 ->
            return (IrLiteral (FloatValue (s_concat [v, "f"])) f)
          _ -> return (IrLiteral (FloatValue v) f)
        (Literal (StringValue s) _) -> return $ IrLiteral (StringValue s) f
        (Literal (BoolValue   b) _) -> return $ IrLiteral (BoolValue b) f
        (This) -> return $ IrPreUnop Deref (IrIdentifier ([], thisPtrName))
        (Self                     ) -> throw $ KitError $ BasicError
          ("unexpected `Self` in typed AST")
          (Just pos)
        (Identifier (Var v)) -> case t of
          TypeTypeOf x _ -> throwk $ BasicError
            "Names of types can't be used as runtime values"
            (Just pos)
          TypeFunction rt args varargs params | not (null params) -> do
            -- generic function
            tctx   <- modTypeContext ctx mod
            params <- mapMWithErrors (follow ctx tctx) params
            return $ IrIdentifier $ monomorphName v params
          _ -> return $ IrIdentifier v
        (StaticMember tp params name) -> do
          tctx   <- modTypeContext ctx mod
          params <- forMWithErrors params $ follow ctx tctx
          when (or $ map typeUnresolved params) $ throwk $ TypingError
            (  s_unpack (showTypePath tp)
            ++ " parameter values for static member couldn't be resolved: "
            ++ show params
            )
            pos
          t <- follow ctx tctx t
          return $ IrIdentifier $ subPath (monomorphName tp params) $ name
        (Identifier (MacroVar v _)) -> return $ IrIdentifier ([], v)
        (TypeAnnotation e1 t      ) -> throw $ KitError $ BasicError
          ("unexpected type annotation in typed AST")
          (Just pos)
        (PreUnop Ref (TypedExpr { tExpr = This })) -> do
          return $ IrIdentifier ([], thisPtrName)
        (PreUnop op e1) -> do
          r1 <- r e1
          return $ IrPreUnop op r1
        (PostUnop op e1) -> do
          r1 <- r e1
          return $ IrPostUnop op r1
        (Binop Mod e1 e2) -> do
          -- FIXME: this special handling is C-specific and should be in
          -- GenerateCode, not GenerateIr
          r1 <- r e1
          r2 <- r e2
          t1 <- findUnderlyingType ctx mod (Just pos) (inferredType e1)
          t2 <- findUnderlyingType ctx mod (Just pos) (inferredType e2)
          let maxFloat = foldr
                (\t acc -> case t of
                  BasicTypeFloat f -> max f acc
                  _                -> acc
                )
                0
                [t1, t2]
          case maxFloat of
            64 -> do
              return $ IrCall (IrIdentifier ([], "fmod")) [r1, r2]
            32 -> do
              return $ IrCall (IrIdentifier ([], "fmodf")) [r1, r2]
            _ -> return $ IrBinop Mod r1 r2
        (Binop Eq (TypedExpr { tExpr = Literal (StringValue s) _ }) e2) -> do
          r2 <- r e2
          return $ stringCompare r2 s
        (Binop Eq e1 (TypedExpr { tExpr = Literal (StringValue s) _ })) -> do
          r1 <- r e1
          return $ stringCompare r1 s
        (Binop Assign (e1@(TypedExpr { tIsLvalue = True, inferredType = TypeArray t n })) (TypedExpr { tExpr = Empty }))
          -> do
            -- assignment of `empty` to arrays should use memset
            r1 <- r e1
            return $ IrCall (IrIdentifier ([], "memset"))
                            [r1, IrIdentifier ([], "0"), IrSizeOf f]
        (Binop Assign e1 (TypedExpr { tExpr = ArrayLiteral values })) -> do
          r1     <- r e1
          values <- mapM r values
          return $ IrCompound
            [ IrBinop Assign
                      (IrArrayAccess r1 (IrLiteral (IntValue i) BasicTypeCSize))
                      val
            | (i, val) <- zip [0 ..] values
            ]
        (Binop op e1 e2) -> do
          r1 <- r e1
          r2 <- r e2
          return $ IrBinop op r1 r2
        (For e1@(TypedExpr { tExpr = Identifier (Var id) }) (TypedExpr { tExpr = RangeLiteral eFrom eTo }) e3)
          -> do
            t     <- findUnderlyingType ctx mod (Just $ tPos e1) (inferredType e1)
            rFrom <- typedToIr ctx (ictx { ictxIsLoopCondition = True }) mod eFrom
            rTo   <- typedToIr ctx (ictx { ictxIsLoopCondition = True }) mod eTo
            r3    <- r e3
            return $ IrFor (tpName id) t rFrom rTo r3
        (While e1 e2 d) -> do
          r1 <- typedToIr ctx (ictx { ictxIsLoopCondition = True }) mod e1
          r2 <- r e2
          return $ IrWhile r1 r2 d
        (If e1 e2 e3) -> do
          r1 <- r e1
          r2 <- r e2
          r3 <- maybeR e3
          return $ IrIf r1 r2 r3
        (Continue ) -> return $ IrContinue
        (Break    ) -> return $ IrBreak
        (Return e1) -> do
          r1 <- maybeR e1
          return $ IrReturn r1
        (Match _ [] Nothing) -> do
          return $ IrBlock []
        (Match e1 cases e2) -> do
          r1        <- r e1
          r2        <- maybeR e2
          matchType <- findUnderlyingType ctx mod (Just pos) (inferredType e1)
          let complexMatch = case matchType of
                BasicTypeComplexEnum _  -> True
                BasicTypeTuple _ _      -> True
                BasicTypeStruct _       -> True
                BasicTypeAnonStruct _ _ -> True
                BasicTypeUnion _        -> True
                BasicTypeAnonUnion _ _  -> True
                _                       -> False
          case matchType of
            _ | complexMatch -> do
              -- complex match with ADT
              cases' <- forMWithErrors cases $ \c -> do
                (conditions, assigns) <- patternMatch ctx
                                                      mod
                                                      r
                                                      (matchPattern c)
                                                      matchType
                                                      r1
                body <- r $ matchBody c
                let mergeConditions x = if null x
                      then (IrLiteral (BoolValue True) BasicTypeBool)
                      else
                        let (h, t) = (head x, tail x)
                        in  if null t
                              then h
                              else (IrBinop And h (mergeConditions t))
                let
                  (fullAssigns, _) = foldr
                    (\(varName, varType, varExpr) (l, n) ->
                      ( let tmpName = (s_concat ["__match__", s_pack $ show n])
                        in
                          l
                            ++ [ IrVarDeclaration tmpName varType (Just varExpr)
                               , IrVarDeclaration
                                 varName
                                 varType
                                 (Just $ IrIdentifier ([], tmpName))
                               ]
                      , n + 1
                      )
                    )
                    ([], 1)
                    assigns
                return
                  $ ( mergeConditions conditions
                    , (IrBlock $ fullAssigns ++ [body])
                    )
              let ifX ((cond, body) : t) =
                    IrIf cond body (if null t then r2 else Just $ ifX t)
              return $ ifX cases'
            BasicTypeBool -> do
              -- transform into an if statement
              let branchOrDefault b = case b of
                    Just x -> do
                      rx <- r $ matchBody x
                      return $ Just rx
                    Nothing -> return r2
              trueBranch <- branchOrDefault $ find
                (\c -> case (tExpr $ matchPattern c) of
                  Literal (BoolValue True) _ -> True
                  _                          -> False
                )
                cases
              falseBranch <- branchOrDefault $ find
                (\c -> case (tExpr $ matchPattern c) of
                  Literal (BoolValue False) _ -> True
                  _                           -> False
                )
                cases
              case (trueBranch, falseBranch) of
                (Just a, Just b) | a == b -> return $ a
                (Just a, b     )          -> return $ IrIf r1 a b
                (Nothing, Just b) ->
                  return $ IrIf (IrPreUnop Invert r1) b Nothing
                (Nothing, Nothing) -> throwk $ InternalError
                  "Boolean match with no true or false branches"
                  (Just pos)
            _ -> do
              -- simple match
              cases' <- forMWithErrors cases $ \c -> do
                pattern <- r $ matchPattern c
                body    <- r $ matchBody c
                return (pattern, body)
              def <- maybeR e2
              let ifX ((a, b) : t) = IrIf
                    (case a of
                      IrLiteral (StringValue s) _ -> stringCompare r1 s
                      _                           -> (IrBinop Eq r1 a)
                    )
                    b
                    (if null t then r2 else Just $ ifX t)
              return $ ifX cases'
  
        (InlineCall e1   ) -> return $ undefined -- TODO
        (Field e1 (Var v)) -> do
          r1 <- r e1
          return $ IrField r1 (tpName v)
        (Field e1 (MacroVar v _)) -> throwk $ InternalError
          ("unexpected macro var (" ++ (s_unpack v) ++ ") in typed AST")
          (Just pos)
        (ArrayAccess e1 e2) -> do
          r1 <- r e1
          r2 <- r e2
          return $ IrArrayAccess r1 r2
        (Call e1 implicits explicitArgs) -> do
          let args = implicits ++ explicitArgs
          r1    <- r e1
          args' <- mapMWithErrors r args
          return $ IrCall r1 args'
        (Cast e1 t2) -> do
          r1  <- r e1
          t1' <- findUnderlyingType ctx mod (Just pos) (inferredType e1)
          return $ if t1' == f then r1 else IrCast r1 f
        (Unsafe       e1   ) -> r e1
        (BlockComment s    ) -> return $ IrBlock []
        (RangeLiteral e1 e2) -> throwk
          $ BasicError ("unexpected range literal in typed AST") (Just pos)
        (ArrayLiteral items) -> do
          items' <- mapMWithErrors r items
          let
            contentType = case f of
              CArray t _                  -> t
              BasicTypeConst (CArray t _) -> t
              _                           -> throwk $ BasicError
                ("unexpected array literal type: " ++ show f)
                (Just pos)
          return $ IrCArrLiteral items' contentType
        (LocalVarDeclaration (Var name) ts const def) -> do
          case (def, f) of
            (Just x, CArray _ _)
              | case x of
                TypedExpr { tExpr = Empty } -> False
                _                           -> True
              -> do
                a <- typedToIr ctx ictx mod $ makeExprTyped
                  (LocalVarDeclaration (Var name) ts const Nothing)
                  (inferredType e)
                  pos
                b <- typedToIr ctx ictx mod $ makeExprTyped
                  (Binop
                    Assign
                    (makeExprTyped (Identifier (Var name)) (inferredType e) pos)
                    x
                  )
                  (inferredType e)
                  pos
                return $ IrCompound [a, b]
            _ -> do
              def <- maybeR def
              return $ IrVarDeclaration (tpName name) f def
        (LocalVarDeclaration (MacroVar v _) _ _ _) -> do
          throwk $ BasicError
            ("unexpected macro var (" ++ (s_unpack v) ++ ") in typed AST")
            (Just pos)
        -- (Defer x) -> do
        --   throwk $ InternalError "Not yet implemented" (Just pos)
        (StructInit t fields) -> do
          resolvedFields <- forMWithErrors
            fields
            (\(name, e1) -> do
              r1 <- r e1
              return (name, r1)
            )
          return $ IrStructInit f resolvedFields
        (UnionInit t (name, e)) -> do
          resolved <- r e
          return $ IrUnionInit f (name, resolved)
        (EnumInit (TypeInstance tp p) discriminant args) -> do
          tctx           <- modTypeContext ctx mod
          resolvedParams <- forM p $ follow ctx tctx
          f              <- findUnderlyingType ctx
                                               mod
                                               (Just pos)
                                               (TypeInstance tp resolvedParams)
          let disc = if null $ tpNamespace discriminant
                then discriminant
                else subPath (monomorphName tp resolvedParams)
                             (tpName discriminant)
          resolvedArgs <- forMWithErrors args $ \(name, x) -> do
            x <- r x
            return (name, x)
          return $ IrEnumInit f disc resolvedArgs
        (EnumField x variantName fieldName) -> do
          r1 <- r x
          let (TypeInstance tp params) = inferredType x
          return
            $ IrField
                (IrField (IrField r1 variantFieldName) $ discriminantMemberName
                  (subPath (monomorphName tp params) variantName)
                )
            $ fieldName
        (TupleInit slots) -> do
          resolvedSlots <- forMWithErrors
            slots
            (\s -> do
              r1 <- r s
              return r1
            )
          return $ IrTupleInit f resolvedSlots
        (TupleSlot x slot) -> do
          r1 <- r x
          return $ IrField r1 (s_pack $ "__slot" ++ show slot)
        (Box i@(TraitImplementation { implTrait = TypeTraitConstraint ((modPath, traitName), params) }) e1)
          -> do
            r1     <- r e1
            for    <- findUnderlyingType ctx mod (Just pos) $ implFor i
            tctx   <- modTypeContext ctx mod
            params <- forMWithErrors params $ follow ctx tctx
            let structName =
                  subPath (monomorphName (modPath, traitName) params) "box"
            return $ IrStructInit
              (BasicTypeStruct structName)
              [ (valuePointerName, r1)
              , ( vtablePointerName
                , IrPreUnop
                  Ref
                  (IrIdentifier $ monomorphName
                    (monomorphName (modPath, traitName) params)
                    [implFor i]
                  )
                )
              ]
        (Box t _) -> throwk $ InternalError
          ("Invalid boxed implementation: " ++ show t)
          (Just pos)
        (BoxedValue x) -> do
          box <- r x
          return $ IrField box valuePointerName
        (BoxedVtable trait x) -> do
          box <- r x
          return $ IrPreUnop Deref (IrField box vtablePointerName)
        (StaticVtable i@(TraitImplementation { implTrait = TypeTraitConstraint ((modPath, traitName), params) }))
          -> do
            return $ IrIdentifier $ monomorphName
              (monomorphName (modPath, traitName) params)
              [implFor i]
  
        (SizeOf t) -> do
          t <- findUnderlyingType ctx mod (Just pos) t
          return $ IrSizeOf t
        (Null ) -> return IrNull
        (Empty) -> do
          t <- findUnderlyingType ctx mod (Just pos) t
          case t of
            CArray              _ _ -> return ()
            BasicTypeAnonStruct _ _ -> return ()
            BasicTypeStruct      _  -> return ()
            BasicTypeComplexEnum _  -> return ()
            _                       -> throwk $ TypingError
              ("`empty` isn't a valid value of type " ++ show t)
              pos
          return $ IrEmpty t
        (Undefined) ->
          throwk $ TypingError "`undefined` isn't a valid runtime value" pos
        (VarArg x) -> do
          return $ IrCall (IrIdentifier ([], "va_arg"))
                          [IrIdentifier ([], x), IrType f]
        (           VarArgListCopy x) -> return $ IrIdentifier ([], x)
        InlineCExpr s               t -> return $ IrInlineC s
        t                             -> do
          throwk $ InternalError
            ("Unexpected expression in typed AST:\n\n" ++ show t)
            (Just pos)
  