MycLib/Src/AST/Myc.Ast.Compiler.Specializer.pas

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 (e.g. internal compilation failures)
    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;
        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;

    TCompiledFunction = record
    public
        Func: TDataValue.TFunc;
        StaticType: IStaticType;
        IsPure: Boolean;
        constructor Create(const AFunc: TDataValue.TFunc; const AStaticType: IStaticType; AIsPure: Boolean);
    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.
    // It propagates Purity information but does NOT perform Constant Folding yet.
    TStaticSpecializer = class(TAstTransformer, IAstSpecializer)
        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;

    protected
        function VisitFunctionCall(const Node: IFunctionCallNode): IAstNode; 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.');
    // CompileFunc can be nil if no recursion is supported, but usually required.

    FMonomorphCache := AMonomorphCache;
    FFunctionRegistry := AFunctionRegistry;
    FCompileFunc := ACompileFunc;
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: 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(Node.Identity, 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(Node.Identity, 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 FMonomorphCache.TryGetFunction(key, specializedMethod) then
    begin
        // 4a. Cache Hit (Environment)
        Result :=
            TAst.FunctionCall(
                Node.Identity,
                newCallee,
                newArgs,
                specializedMethod.ReturnType,
                Node.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,
                Node.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 := TAst.FunctionCall(Node.Identity, newCallee, newArgs, Node.StaticType, Node.IsTailCall, nil);
                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.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, Node.IsTailCall, compiled.Func, compiled.IsPure);
        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(Node.Identity, newCallee, newArgs, Node.StaticType, Node.IsTailCall, nil);
end;

{ TMonoCacheKey }

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

{ TCompiledFunction }

constructor TCompiledFunction.Create(const AFunc: TDataValue.TFunc; const AStaticType: IStaticType; AIsPure: Boolean);
begin
    Func := AFunc;
    StaticType := AStaticType;
    IsPure := AIsPure;
end;

end.