unit Myc.Ast.Compiler.Specializer;

interface

uses
    System.SysUtils,
    System.Generics.Collections,
    Myc.Data.Value,
    Myc.Ast,
    Myc.Ast.Types,
    Myc.Ast.Nodes,
    Myc.Ast.Visitor,
    Myc.Ast.Scope,
    Myc.Ast.Environment,
    Myc.Ast.RTL;

type
    IAstSpecializer = interface(IAstVisitor)
        function Execute(const RootNode: IAstNode): IAstNode;
    end;

    // This transformer runs *after* TypeChecker
    // It specializes all statically resolvable function calls (RTL and user-defined)
    // by replacing them with nodes that have a direct StaticTarget.
    TStaticSpecializer = class(TAstTransformer, IAstSpecializer)
    private
        FEnvironment: IEnvironment;

        function GetStaticRtlFunction(const AName: string; const AArgTypes: TArray<IStaticType>): TSpecializedMethod;

    protected
        function VisitFunctionCall(const Node: IFunctionCallNode): IAstNode; override;

    public
        constructor Create(const AEnvironment: IEnvironment);
        function Execute(const RootNode: IAstNode): IAstNode;

        class function Specialize(
            const AEnvironment: IEnvironment;
            const RootNode: IAstNode;
            const ADescriptor: IScopeDescriptor
        ): IAstNode; static;
    end;

implementation

{ TStaticSpecializer }

constructor TStaticSpecializer.Create(const AEnvironment: IEnvironment);
begin
    inherited Create;
    Assert(Assigned(AEnvironment));
    FEnvironment := AEnvironment;
end;

class function TStaticSpecializer.Specialize(
    const AEnvironment: IEnvironment;
    const RootNode: IAstNode;
    const ADescriptor: IScopeDescriptor
): IAstNode;
begin
    // ADescriptor is no longer needed here, the environment has all it needs
    var specializer := TStaticSpecializer.Create(AEnvironment) as IAstSpecializer;
    Result := specializer.Execute(RootNode);
end;

function TStaticSpecializer.Execute(const RootNode: IAstNode): IAstNode;
begin
    Result := Accept(RootNode);
    if not Assigned(Result) then
        Result := TAst.Block([]);
end;

function TStaticSpecializer.GetStaticRtlFunction(const AName: string; const AArgTypes: TArray<IStaticType>): TSpecializedMethod;
begin
    // Helper to query the RTL bootstrap cache
    Result := TRtlRegistry.GetStaticSpecialization(AName, AArgTypes);
end;

function TStaticSpecializer.VisitFunctionCall(const Node: IFunctionCallNode): IAstNode;
var
    newCallee: IAstNode;
    newArgs: TArray<IAstNode>;
    i: Integer;
    calleeIdent: IIdentifierNode;
    argTypes: TArray<IStaticType>;
    allTypesKnown: Boolean;
    funcName: string;
    key: TMonoCacheKey;
    specializedMethod: TSpecializedMethod;
    funcDef: IFunctionDefinition;
begin
    // 1. Specialize children first (bottom-up)
    newCallee := Accept(Node.Callee);
    newArgs := AcceptNodes(Node.Arguments);

    // 2. Check if this call is a candidate for specialization
    if newCallee.Kind <> akIdentifier then
    begin
        if (newCallee = Node.Callee) and (newArgs = Node.Arguments) then
            Result := Node
        else
            Result := TAst.FunctionCall(newCallee, newArgs, Node.StaticType, Node.IsTailCall, nil);
        exit;
    end;

    calleeIdent := newCallee.AsIdentifier;
    funcName := calleeIdent.Name;

    // 3. Check if all argument types are statically known
    allTypesKnown := True;
    SetLength(argTypes, Length(newArgs));
    for i := 0 to High(newArgs) do
    begin
        argTypes[i] := newArgs[i].AsTypedNode.StaticType;
        if argTypes[i].Kind = stUnknown then
        begin
            allTypesKnown := False;
            break;
        end;
    end;

    if not allTypesKnown then
    begin
        if (newCallee = Node.Callee) and (newArgs = Node.Arguments) then
            Result := Node
        else
            Result := TAst.FunctionCall(newCallee, newArgs, Node.StaticType, Node.IsTailCall, nil);
        exit;
    end;

    // --- At this point, the call is statically resolvable ---

    // 4. Check the Environment (Instance) Cache
    key := TMonoCacheKey.Create(calleeIdent.Address, argTypes);

    if FEnvironment.MonomorphCache.TryGetValue(key, specializedMethod) then
    begin
        // 4a. Cache Hit (Environment): Found user-defined or previously specialized fn
        Result :=
            TAst.FunctionCall(
                newCallee,
                newArgs,
                specializedMethod.ReturnType, // Use the type from the cache
                Node.IsTailCall,
                specializedMethod.Target // Use the TFunc from the cache
            );
        exit;
    end;

    // 5. Check the RTL (Global) Bootstrap Cache
    specializedMethod := GetStaticRtlFunction(funcName, argTypes);
    if Assigned(specializedMethod.Target) then
    begin
        // 5a. Cache Hit (RTL): Found a native, static RTL function
        FEnvironment.MonomorphCache.Add(key, specializedMethod);

        Result :=
            TAst.FunctionCall(
                newCallee,
                newArgs,
                specializedMethod.ReturnType, // Use the type from the cache
                Node.IsTailCall,
                specializedMethod.Target // Use the TFunc from the cache
            );
        exit;
    end;

    // 6. Cache Miss (User Code)
    funcDef := FEnvironment.FunctionRegistry.Resolve(calleeIdent.Address);

    if (funcDef <> nil) then
    begin
        // Cannot specialize closures or functions with nested closures
        // that might depend on the local environment being detached.
        // We check if it's a lambda expression node (which it usually is)
        if funcDef.Kind = akLambdaExpression then
        begin
            var lambdaDef := funcDef.AsLambdaExpression;
            if (Length(lambdaDef.Upvalues) > 0) or (lambdaDef.HasNestedLambdas) then
            begin
                // This function relies on local scope context. We cannot lift it
                // to a static compilation unit safely. Fallback to dynamic dispatch.
                Result := TAst.FunctionCall(newCallee, newArgs, Node.StaticType, Node.IsTailCall, nil);
                exit;
            end;
        end;

        // 6a. Compile func with KNOWN TYPES
        var compiled := FEnvironment.Compile(funcDef, argTypes);

        // 6b. Store in cache
        // Get the return type from the *full function type* returned by Compile
        var returnType := compiled.StaticType.Signatures[0].ReturnType;
        specializedMethod := TSpecializedMethod.Create(compiled.Func, returnType);
        FEnvironment.MonomorphCache.Add(key, specializedMethod);

        // 6e. Return the new node
        Result :=
            TAst.FunctionCall(
                newCallee,
                newArgs,
                returnType, // Use the extracted return type
                Node.IsTailCall,
                compiled.Func // Use the compiled TFunc
            );
        exit;
    end;

    // 7. Fallback: Not RTL, Not User-Code (oder AST nicht gefunden) -> Dynamic
    if (newCallee = Node.Callee) and (newArgs = Node.Arguments) then
        Result := Node
    else
        Result := TAst.FunctionCall(newCallee, newArgs, Node.StaticType, Node.IsTailCall, nil);
end;

end.