program ClassExtension;

{$APPTYPE CONSOLE}

{$R *.res}

uses
  Generics.Collections,
  SysUtils,
  Spring.VirtualClass;

type
  TFoo = class

  end;

  TFooBar = class(TFoo)

  end;

  TRemoveInstanceProc = procedure(const key: TObject) of object;

  TObjectExtension = class helper for TObject
  strict private
    class var proxyClasses: TVirtualClasses;
    class var removeProcs: TList<TRemoveInstanceProc>;
    class var fields: TList<TObject>;
    class procedure FreeInstanceHook(const Self: TObject); static;
    class constructor Create;
    class destructor Destroy;
    class procedure HandleInstanceAdd(Sender: TObject; const Item: TObject;
      Action: TCollectionNotification);
  protected
    class function AddExtension<T>: TDictionary<TObject,T>; static;
  end;

  TFooHelper = class helper(TObjectExtension) for TFoo
  strict private
    class constructor Create;

    class var fBar: TDictionary<TObject,Integer>;
    function GetBar: Integer;
    procedure SetBar(const Value: Integer);
  public
    property Bar: Integer read GetBar write SetBar;
  end;

{ TObjectExtension }

class constructor TObjectExtension.Create;
begin
  proxyClasses := TVirtualClasses.Create;
  removeProcs := TList<TRemoveInstanceProc>.Create;
  fields := TObjectList<TObject>.Create;
end;

class destructor TObjectExtension.Destroy;
begin
  fields.Free;
  removeProcs.Free;
  proxyClasses.Free;
end;

class function TObjectExtension.AddExtension<T>: TDictionary<TObject, T>;
begin
  Result := TDictionary<TObject, T>.Create;
  Result.OnKeyNotify := HandleInstanceAdd;
  fields.Add(Result);
  removeProcs.Add(Result.Remove);
end;

class procedure TObjectExtension.FreeInstanceHook(const Self: TObject);
var
  freeInstance: TFreeInstance;
  proc: TRemoveInstanceProc;
begin
  freeInstance := GetClassData(Self.ClassParent).FreeInstance;
  for proc in removeProcs do
    proc(Self);
  freeInstance(Self);
end;

class procedure TObjectExtension.HandleInstanceAdd(Sender: TObject;
  const Item: TObject; Action: TCollectionNotification);
var
  proxyClass: TClass;
begin
  if Action = cnAdded then
  begin
    proxyClass := proxyClasses.GetVirtualClass(Item.ClassType);
    GetClassData(proxyClass).FreeInstance := FreeInstanceHook;
    proxyClasses.Proxify(Item);
  end;
end;

{ TFooHelper }

class constructor TFooHelper.Create;
begin
  fBar := AddExtension<Integer>;
end;

function TFooHelper.GetBar: Integer;
begin
  fBar.TryGetValue(Self, Result);
end;

procedure TFooHelper.SetBar(const Value: Integer);
begin
  fBar.AddOrSetValue(Self, Value);
end;

var
  f: TFoo;
  b: TFooBar;
begin
  f := TFoo.Create;
  b := TFooBar.Create;
  Writeln(b.Bar);
  b.Bar := 56;
  Writeln(b.Bar);

  Writeln(f.Bar);
  f.Bar := 42;
  f.Bar := 42;
  Writeln(f.Bar);
  f.Free;
  b.Free;
  Readln;
  ReportMemoryLeaksOnShutdown := True;
end.