<commit_before>module Concretize where

import Control.Monad.State
import qualified Data.Map as Map
import Data.Maybe (fromMaybe)
import qualified Data.Set as Set
import Data.Set ((\\))
import Data.List (foldl')
import Debug.Trace

import Obj
import Constraints
import Types
import Util
import TypeError
import AssignTypes
import Polymorphism
import InitialTypes
import Lookup

-- | This function performs two things:
-- |  1. Finds out which polymorphic functions that needs to be added to the environment for the calls in the function to work.
-- |  2. Changes the name of symbols at call sites so they use the polymorphic name
-- |  Both of these results are returned in a tuple: (<new xobj>, <dependencies>)
concretizeXObj :: Bool -> TypeEnv -> Env -> [SymPath] -> XObj -> Either TypeError (XObj, [XObj])
concretizeXObj allowAmbiguityRoot typeEnv rootEnv visitedDefinitions root =
  case runState (visit allowAmbiguityRoot rootEnv root) [] of
    (Left err, _) -> Left err
    (Right xobj, deps) -> Right (xobj, deps)
  where
    rootDefinitionName :: String
    rootDefinitionName = getName root

    visit :: Bool -> Env -> XObj -> State [XObj] (Either TypeError XObj)
    visit allowAmbig env xobj@(XObj (Sym _ _) _ _) = visitSymbol allowAmbig env xobj
    visit allowAmbig env xobj@(XObj (MultiSym _ _) _ _) = visitMultiSym allowAmbig env xobj
    visit allowAmbig env xobj@(XObj (InterfaceSym _) _ _) = visitInterfaceSym allowAmbig env xobj
    visit allowAmbig env xobj@(XObj (Lst _) i t) =
      do visited <- visitList allowAmbig env xobj
         return $ do okVisited <- visited
                     Right (XObj (Lst okVisited) i t)
    visit allowAmbig env xobj@(XObj (Arr arr) i (Just t)) =
      do visited <- fmap sequence (mapM (visit allowAmbig env) arr)
         concretizeTypeOfXObj typeEnv xobj
         return $ do okVisited <- visited
                     Right (XObj (Arr okVisited) i (Just t))
    visit _ _ x = return (Right x)

    visitList :: Bool -> Env -> XObj -> State [XObj] (Either TypeError [XObj])
    visitList _ _ (XObj (Lst []) _ _) = return (Right [])

    visitList _ env (XObj (Lst [defn@(XObj Defn _ _), nameSymbol@(XObj (Sym (SymPath [] "main") _) _ _), args@(XObj (Arr argsArr) _ _), body]) _ _) =
      if not (null argsArr)
      then return $ Left (MainCannotHaveArguments nameSymbol (length argsArr))
      else do visitedBody <- visit False env body -- allowAmbig == 'False'
              return $ do okBody <- visitedBody
                          let t = fromMaybe UnitTy (ty okBody)
                          if not (isTypeGeneric t) && t /= UnitTy && t /= IntTy
                            then Left (MainCanOnlyReturnUnitOrInt nameSymbol t)
                            else return [defn, nameSymbol, args, okBody]

    visitList _ env (XObj (Lst [defn@(XObj Defn _ _), nameSymbol, args@(XObj (Arr argsArr) _ _), body]) _ t) =
      do mapM_ (concretizeTypeOfXObj typeEnv) argsArr
         let functionEnv = Env Map.empty (Just env) Nothing [] InternalEnv 0
             envWithArgs = foldl' (\e arg@(XObj (Sym (SymPath _ argSymName) _) _ _) ->
                                     extendEnv e argSymName arg)
                                  functionEnv argsArr
             Just funcTy = t
             allowAmbig = isTypeGeneric funcTy
         visitedBody <- visit allowAmbig envWithArgs body
         return $ do okBody <- visitedBody
                     return [defn, nameSymbol, args, okBody]

    -- | Fn / λ
    visitList allowAmbig env (XObj (Lst [(XObj (Fn _ _) fni fnt), args@(XObj (Arr argsArr) ai at), body]) i t) =
      -- The basic idea of this function is to first visit the body of the lambda ("in place"),
      -- then take the resulting body and put into a separate function 'defn' with a new name
      -- in the global scope. That function definition will be set as the lambdas '.callback' in
      -- the C code.
      do let Just ii = i
             Just funcTy = t
              -- | TODO: This code is a copy of the one above in Defn, remove duplication:
             functionEnv = Env Map.empty (Just env) Nothing [] InternalEnv (envFunctionNestingLevel env + 1)
             envWithArgs = foldl' (\e arg@(XObj (Sym (SymPath _ argSymName) _) _ _) ->
                                     extendEnv e argSymName arg)
                                  functionEnv argsArr
         visitedBody <- visit allowAmbig envWithArgs body
         case visitedBody of
           Right (okBody) ->
             let capturedVars = collectCapturedVars okBody
                 structMemberPairs = concatMap (\(XObj (Sym path _) _ (Just symTy)) ->
                                                  [XObj (Sym path Symbol) Nothing Nothing, tyToXObj symTy])
                                     capturedVars

                 lambdaName = "_Lambda_" ++ rootDefinitionName ++ "_" ++ show (infoIdentifier ii)
                 lambdaNameSymbol = XObj (Sym (SymPath [] lambdaName) Symbol) (Just dummyInfo) Nothing
                 extendedArgs = if null capturedVars
                                then args
                                     -- If the lambda captures anything it need an extra arg for its env:
                                else XObj (Arr (XObj (Sym (SymPath [] "_env") Symbol)
                                                (Just dummyInfo)
                                                (Just (PointerTy (StructTy environmentTypeName []))) :
                                                argsArr)) ai at
                 lambdaCallback = XObj (Lst [XObj Defn (Just dummyInfo) Nothing, lambdaNameSymbol, extendedArgs, okBody]) i t

                 environmentTypeName = lambdaName ++ "_env"
                 environmentStructTy = StructTy environmentTypeName []
                 environmentStruct = XObj (Lst (XObj (Typ environmentStructTy) Nothing Nothing :
                                                XObj (Sym (SymPath [] environmentTypeName) Symbol) Nothing Nothing :
                                                XObj (Arr structMemberPairs) Nothing Nothing :
                                                [])) i (Just TypeTy)

                 -- The type env has to contain the lambdas environment struct for 'concretizeDefinition' to work:
                 extendedTypeEnv = TypeEnv (extendEnv (getTypeEnv typeEnv) environmentTypeName environmentStruct)

             in case concretizeDefinition allowAmbig extendedTypeEnv env visitedDefinitions lambdaCallback funcTy of
                  Left err -> return (Left err)
                  Right (concreteLiftedLambda, deps) ->
                    do modify (environmentStruct :)
                       modify (concreteLiftedLambda :)
                       modify (deps ++)
                       return (Right [XObj (Fn (Just lambdaName) (Set.fromList capturedVars)) fni fnt, args, okBody])
           _ ->
             error "Visited body isn't a defn."

    visitList _ env (XObj (Lst [def@(XObj Def _ _), nameSymbol, body]) _ t) =
      do let Just defTy = t
             allowAmbig = isTypeGeneric defTy<commit_msg>Comments.<commit_after>module Concretize where

import Control.Monad.State
import qualified Data.Map as Map
import Data.Maybe (fromMaybe)
import qualified Data.Set as Set
import Data.Set ((\\))
import Data.List (foldl')
import Debug.Trace

import Obj
import Constraints
import Types
import Util
import TypeError
import AssignTypes
import Polymorphism
import InitialTypes
import Lookup

-- | This function performs two things:
-- |  1. Finds out which polymorphic functions that needs to be added to the environment for the calls in the function to work.
-- |  2. Changes the name of symbols at call sites so they use the polymorphic name
-- |  Both of these results are returned in a tuple: (<new xobj>, <dependencies>)
concretizeXObj :: Bool -> TypeEnv -> Env -> [SymPath] -> XObj -> Either TypeError (XObj, [XObj])
concretizeXObj allowAmbiguityRoot typeEnv rootEnv visitedDefinitions root =
  case runState (visit allowAmbiguityRoot rootEnv root) [] of
    (Left err, _) -> Left err
    (Right xobj, deps) -> Right (xobj, deps)
  where
    rootDefinitionName :: String
    rootDefinitionName = getName root

    visit :: Bool -> Env -> XObj -> State [XObj] (Either TypeError XObj)
    visit allowAmbig env xobj@(XObj (Sym _ _) _ _) = visitSymbol allowAmbig env xobj
    visit allowAmbig env xobj@(XObj (MultiSym _ _) _ _) = visitMultiSym allowAmbig env xobj
    visit allowAmbig env xobj@(XObj (InterfaceSym _) _ _) = visitInterfaceSym allowAmbig env xobj
    visit allowAmbig env xobj@(XObj (Lst _) i t) =
      do visited <- visitList allowAmbig env xobj
         return $ do okVisited <- visited
                     Right (XObj (Lst okVisited) i t)
    visit allowAmbig env xobj@(XObj (Arr arr) i (Just t)) =
      do visited <- fmap sequence (mapM (visit allowAmbig env) arr)
         concretizeTypeOfXObj typeEnv xobj
         return $ do okVisited <- visited
                     Right (XObj (Arr okVisited) i (Just t))
    visit _ _ x = return (Right x)

    visitList :: Bool -> Env -> XObj -> State [XObj] (Either TypeError [XObj])
    visitList _ _ (XObj (Lst []) _ _) = return (Right [])

    visitList _ env (XObj (Lst [defn@(XObj Defn _ _), nameSymbol@(XObj (Sym (SymPath [] "main") _) _ _), args@(XObj (Arr argsArr) _ _), body]) _ _) =
      if not (null argsArr)
      then return $ Left (MainCannotHaveArguments nameSymbol (length argsArr))
      else do visitedBody <- visit False env body -- allowAmbig == 'False'
              return $ do okBody <- visitedBody
                          let t = fromMaybe UnitTy (ty okBody)
                          if not (isTypeGeneric t) && t /= UnitTy && t /= IntTy
                            then Left (MainCanOnlyReturnUnitOrInt nameSymbol t)
                            else return [defn, nameSymbol, args, okBody]

    visitList _ env (XObj (Lst [defn@(XObj Defn _ _), nameSymbol, args@(XObj (Arr argsArr) _ _), body]) _ t) =
      do mapM_ (concretizeTypeOfXObj typeEnv) argsArr
         let functionEnv = Env Map.empty (Just env) Nothing [] InternalEnv 0
             envWithArgs = foldl' (\e arg@(XObj (Sym (SymPath _ argSymName) _) _ _) ->
                                     extendEnv e argSymName arg)
                                  functionEnv argsArr
             Just funcTy = t
             allowAmbig = isTypeGeneric funcTy
         visitedBody <- visit allowAmbig envWithArgs body
         return $ do okBody <- visitedBody
                     return [defn, nameSymbol, args, okBody]

    -- | Fn / λ
    visitList allowAmbig env (XObj (Lst [(XObj (Fn _ _) fni fnt), args@(XObj (Arr argsArr) ai at), body]) i t) =
      -- The basic idea of this function is to first visit the body of the lambda ("in place"),
      -- then take the resulting body and put into a separate function 'defn' with a new name
      -- in the global scope. That function definition will be set as the lambdas '.callback' in
      -- the C code.
      do let Just ii = i
             Just funcTy = t
              -- | TODO: This code is a copy of the one above in Defn, remove duplication:
             functionEnv = Env Map.empty (Just env) Nothing [] InternalEnv (envFunctionNestingLevel env + 1)
             envWithArgs = foldl' (\e arg@(XObj (Sym (SymPath _ argSymName) _) _ _) ->
                                     extendEnv e argSymName arg)
                                  functionEnv argsArr
         visitedBody <- visit allowAmbig envWithArgs body
         case visitedBody of
           Right (okBody) ->
             let -- Analyse the body of the lambda to find what variables it captures
                 capturedVars = collectCapturedVars okBody

                 -- Create a new (top-level) function that will be used when the lambda is called.
                 -- Its name will contain the name of the (normal, non-lambda) function it's contained within,
                 -- plus the identifier of the particular s-expression that defines the lambda.
                 lambdaName = "_Lambda_" ++ rootDefinitionName ++ "_" ++ show (infoIdentifier ii)
                 lambdaNameSymbol = XObj (Sym (SymPath [] lambdaName) Symbol) (Just dummyInfo) Nothing
                 extendedArgs = if null capturedVars
                                then args
                                     -- If the lambda captures anything it need an extra arg for its env:
                                else XObj (Arr (XObj (Sym (SymPath [] "_env") Symbol)
                                                (Just dummyInfo)
                                                (Just (PointerTy (StructTy environmentTypeName []))) :
                                                argsArr)) ai at
                 lambdaCallback = XObj (Lst [XObj Defn (Just dummyInfo) Nothing, lambdaNameSymbol, extendedArgs, okBody]) i t

                 -- The lambda will also carry with it a special made struct containing the variables it captures
                 -- (if it captures at least one variable)
                 structMemberPairs = concatMap (\(XObj (Sym path _) _ (Just symTy)) ->
                                                  [XObj (Sym path Symbol) Nothing Nothing, tyToXObj symTy])
                                     capturedVars
                 environmentTypeName = lambdaName ++ "_env"
                 environmentStructTy = StructTy environmentTypeName []
                 environmentStruct = XObj (Lst (XObj (Typ environmentStructTy) Nothing Nothing :
                                                XObj (Sym (SymPath [] environmentTypeName) Symbol) Nothing Nothing :
                                                XObj (Arr structMemberPairs) Nothing Nothing :
                                                [])) i (Just TypeTy)

                 -- The type env has to contain the lambdas environment struct for 'concretizeDefinition' to work:
                 extendedTypeEnv = TypeEnv (extendEnv (getTypeEnv typeEnv) environmentTypeName environmentStruct)

             in case concretizeDefinition allowAmbig extendedTypeEnv env visitedDefinitions lambdaCallback funcTy of
                  Left err -> return (Left err)
                  Right (concreteLiftedLambda, deps) ->
                    do modify (environmentStruct :)
                       modify (concreteLiftedLambda :)
                       modify (deps ++)
                       return (Right [XObj (Fn (Just lambdaName) (Set.fromList capturedVars)) fni fnt, args, okBody])
           _ ->
             error "Visited body isn't a defn."

    visitList _ env (XObj (Lst [def@(XObj Def _ _), nameSymbol, body]) _ t) =
      do let Just defTy = t
             allowAmbig = isTypeGeneric defTy
{'fraction_non_alphanumeric': 0.07644220968947475, 'fraction_numerical': 0.0005435890466807093, 'mean_word_length': 1.8943952802359882, 'pattern_counts': {'":': 0, '<': 25, '<?xml version=': 0, '>': 47, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 46, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '5036000', 'n_tokens_mistral': 3897, 'n_tokens_neox': 3749, 'n_words': 1447}