unit Myc.Ast.Compiler.Specializer;

interface

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

type
    // Exception specific to specialization errors
    ESpecializerException = class(EAstException);

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

    // --- Monomorphization Cache Definitions ---

    TMonoCacheKey = record
    public
        Address: TResolvedAddress;
        ArgTypes: TArray<IStaticType>;
        Func: TDataValue.TFunc; // Usually nil in key, but part of structure if needed
        constructor Create(const AAddress: TResolvedAddress; const AArgTypes: TArray<IStaticType>);
    end;

    IMonomorphCache = interface
        function TryGetFunction(const Key: TMonoCacheKey; out Func: TSpecializedMethod): Boolean;
        procedure Add(const Key: TMonoCacheKey; const Func: TSpecializedMethod);
    end;

    // This transformer runs *after* TypeChecker.
    // It specializes all statically resolvable function calls.
    TStaticSpecializer = class(TAstTransformer, IAstSpecializer)
    public
        type
            TCompileFunc = reference to function(const Node: IFunctionDefinition; const ArgTypes: TArray<IStaticType>): TCompiledFunction;
    private
        FMonomorphCache: IMonomorphCache;
        FFunctionRegistry: IFunctionDefinitionRegistry;
        FCompileFunc: TCompileFunc;

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

    strict private
        // Specialization Handler (IAstNode signature)
        function VisitFunctionCall(const Node: IAstNode): IAstNode;

    protected
        procedure SetupHandlers; override;

    public
        constructor Create(
            const AMonomorphCache: IMonomorphCache;
            const AFunctionRegistry: IFunctionDefinitionRegistry;
            const ACompileFunc: TCompileFunc
        );
        function Execute(const RootNode: IAstNode): IAstNode;

        class function Specialize(
            const RootNode: IAstNode;
            const AMonomorphCache: IMonomorphCache;
            const AFunctionRegistry: IFunctionDefinitionRegistry;
            const ACompileFunc: TCompileFunc
        ): IAstNode; static;
    end;

implementation

uses
    System.Hash;

{ TStaticSpecializer }

constructor TStaticSpecializer.Create(
    const AMonomorphCache: IMonomorphCache;
    const AFunctionRegistry: IFunctionDefinitionRegistry;
    const ACompileFunc: TCompileFunc
);
begin
    inherited Create;
    if not Assigned(AMonomorphCache) then
        raise ESpecializerException.Create('MonomorphCache cannot be nil.');
    if not Assigned(AFunctionRegistry) then
        raise ESpecializerException.Create('FunctionRegistry cannot be nil.');

    FMonomorphCache := AMonomorphCache;
    FFunctionRegistry := AFunctionRegistry;
    FCompileFunc := ACompileFunc;
end;

procedure TStaticSpecializer.SetupHandlers;
begin
    inherited SetupHandlers; // Load default transformations
    // Override FunctionCall logic
    Register(akFunctionCall, VisitFunctionCall);
end;

class function TStaticSpecializer.Specialize(
    const RootNode: IAstNode;
    const AMonomorphCache: IMonomorphCache;
    const AFunctionRegistry: IFunctionDefinitionRegistry;
    const ACompileFunc: TCompileFunc
): IAstNode;
begin
    var specializer := TStaticSpecializer.Create(AMonomorphCache, AFunctionRegistry, ACompileFunc) 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([], nil);
end;

function TStaticSpecializer.GetStaticRtlFunction(const AName: string; const AArgTypes: TArray<IStaticType>): TSpecializedMethod;
begin
    Result := TRtlRegistry.GetStaticSpecialization(AName, AArgTypes);
end;

function TStaticSpecializer.VisitFunctionCall(const Node: IAstNode): IAstNode;
var
    C: IFunctionCallNode;
    newCall: IFunctionCallNode;
    newCallee: IAstNode;
    newArgs: ITupleNode;
    i: Integer;
    calleeIdent: IIdentifierNode;
    argTypes: TArray<IStaticType>;
    allTypesKnown: Boolean;
    funcName: string;
    key: TMonoCacheKey;
    specializedMethod: TSpecializedMethod;
    funcDef: IFunctionDefinition;
begin
    C := Node.AsFunctionCall;

    // 1. Specialize children first (bottom-up) by calling inherited
    // inherited VisitFunctionCall returns an IAstNode (which is a new IFunctionCallNode if changed)
    newCall := inherited VisitFunctionCall(Node).AsFunctionCall;
    newCallee := newCall.Callee;
    newArgs := newCall.Arguments;

    // 2. Check if this call is a candidate for specialization
    if newCallee.Kind <> akIdentifier then
    begin
        Result := newCall;
        exit;
    end;

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

    // 3. Check if all argument types are statically known
    // Optimization: Use Elements array directly
    var argsElements := newArgs.Elements;
    allTypesKnown := True;
    SetLength(argTypes, Length(argsElements));

    for i := 0 to High(argsElements) do
    begin
        argTypes[i] := argsElements[i].AsTypedNode.StaticType;
        if argTypes[i].Kind = stUnknown then
        begin
            allTypesKnown := False;
            break;
        end;
    end;

    if not allTypesKnown then
    begin
        Result := newCall;
        exit;
    end;

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

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

    if FMonomorphCache.TryGetFunction(key, specializedMethod) then
    begin
        // 4a. Cache Hit (Environment)
        Result :=
            TAst.FunctionCall(
                Node.Identity,
                newCallee,
                newArgs,
                specializedMethod.ReturnType,
                C.IsTailCall,
                specializedMethod.Target,
                specializedMethod.IsPure
            );
        exit;
    end;

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

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

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

    if (funcDef <> nil) then
    begin
        // Cannot specialize closures safely without more complex analysis if they have state
        if funcDef.Kind = akLambdaExpression then
        begin
            var lambdaDef := funcDef.AsLambdaExpression;
            if (Length(lambdaDef.Upvalues) > 0) or (lambdaDef.HasNestedLambdas) then
            begin
                Result := newCall;
                exit;
            end;
        end;

        // 6a. Compile func with KNOWN TYPES
        if not Assigned(FCompileFunc) then
            raise ESpecializerException.Create('Cannot specialize user function: Compiler callback is missing.');

        var compiled := FCompileFunc(funcDef, argTypes);

        // Safety check: The recursive compiler MUST produce a valid function pointer
        if not Assigned(compiled.Func) then
            raise ESpecializerException.CreateFmt('Internal Error: Failed to compile specialization for "%s"', [funcName]);

        // 6b. Store in cache
        var returnType := compiled.StaticType.AsMethod.Signatures[0].ReturnType;
        specializedMethod := TSpecializedMethod.Create(compiled.Func, returnType, compiled.IsPure);
        FMonomorphCache.Add(key, specializedMethod);

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

    // 7. Fallback: Not RTL, Not User-Code -> Dynamic
    Result := newCall;
end;

{ TMonoCacheKey }

constructor TMonoCacheKey.Create(const AAddress: TResolvedAddress; const AArgTypes: TArray<IStaticType>);
begin
    Address := AAddress;
    ArgTypes := AArgTypes;
    Func := nil;
end;

end.