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;

        // Descriptor removed from signature as it is not needed for specialization
        class function Specialize(const AEnvironment: IEnvironment; const RootNode: IAstNode): 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): IAstNode;
begin
    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
    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)
        Result := TAst.FunctionCall(newCallee, newArgs, specializedMethod.ReturnType, Node.IsTailCall, specializedMethod.Target);
        exit;
    end;

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

        Result := TAst.FunctionCall(newCallee, newArgs, specializedMethod.ReturnType, Node.IsTailCall, specializedMethod.Target);
        exit;
    end;

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

    if (funcDef <> nil) then
    begin
        // Cannot specialize closures safely without more complex analysis
        if funcDef.Kind = akLambdaExpression then
        begin
            var lambdaDef := funcDef.AsLambdaExpression;
            if (Length(lambdaDef.Upvalues) > 0) or (lambdaDef.HasNestedLambdas) then
            begin
                Result := TAst.FunctionCall(newCallee, newArgs, Node.StaticType, Node.IsTailCall, nil);
                exit;
            end;
        end;

        // 6a. Compile func with KNOWN TYPES
        // This recursively triggers Bind -> Check -> Specialize for the callee body!
        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);

        // 6c. Return the new node
        Result := TAst.FunctionCall(newCallee, newArgs, returnType, Node.IsTailCall, compiled.Func);
        exit;
    end;

    // 7. Fallback: Not RTL, Not User-Code -> 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.