Refcounting

Cyberboss · Dec 28, 2023 · 616efa3 · 616efa3
1 parent 3ccb916
commit 616efa3
Show file tree

Hide file tree

Showing 10 changed files with 175 additions and 16 deletions.
diff --git a/DMCompiler/DMStandard/_Standard.dm b/DMCompiler/DMStandard/_Standard.dm
@@ -72,6 +72,7 @@ proc/rand(L, H)
 proc/rand_seed(Seed)
 proc/range(Dist, Center)
 proc/ref(Object)
+proc/refcount(Object)
 proc/replacetext(Haystack, Needle, Replacement, Start = 1, End = 0)
 proc/replacetextEx(Haystack, Needle, Replacement, Start = 1, End = 0)
 proc/rgb(R, G, B, A)
@@ -219,9 +220,3 @@ proc/lentext(T)
 
 proc/winshow(player, window, show=1)
 	winset(player, window, "is-visible=[show ? "true" : "false"]")
-
-proc/refcount(var/Object)
-	// woah that's a lot of refs
-	// i wonder if it's true??
-	return 100
-	// (it's not)
diff --git a/OpenDreamRuntime/DreamValue.cs b/OpenDreamRuntime/DreamValue.cs
@@ -436,6 +436,23 @@ public override int GetHashCode() {
         public static bool operator !=(DreamValue a, DreamValue b) {
             return !a.Equals(b);
         }
+
+        [MethodImpl(MethodImplOptions.AggressiveInlining)]
+        public void IncrementDreamObjectRefCount() {
+            if (Type == DreamValueType.DreamObject && _refValue != null) {
+                DreamObject obj = (DreamObject)_refValue;
+                obj.IncrementRefCount();
+            }
+        }
+
+
+        [MethodImpl(MethodImplOptions.AggressiveInlining)]
+        public void DecrementDreamObjectRefCount() {
+            if (Type == DreamValueType.DreamObject && _refValue != null) {
+                DreamObject obj = (DreamObject)_refValue;
+                obj.DecrementRefCount();
+            }
+        }
     }
 
     #region Serialization

diff --git a/OpenDreamRuntime/Objects/DreamObject.cs b/OpenDreamRuntime/Objects/DreamObject.cs
@@ -12,6 +12,8 @@
 using Robust.Shared.Map;
 using Robust.Shared.Serialization.Manager;
 using Robust.Shared.Utility;
+using System.Diagnostics;
+using OpenDreamRuntime.Procs.Native;
 
 namespace OpenDreamRuntime.Objects {
     [Virtual]
@@ -21,6 +23,9 @@ public class DreamObject {
         [Access(typeof(DreamObject))]
         public bool Deleted;
 
+        [Access(typeof(DreamObject), typeof(DreamObjectExtensions), typeof(DreamProcNativeRoot))]
+        public ulong RefCount;
+
         public virtual bool ShouldCallNew => true;
 
         // Shortcuts to IoC dependencies & entity systems
@@ -97,8 +102,14 @@ protected virtual void HandleDeletion() {
 
             Tag = null;
             Deleted = true;
+
             //we release all relevant information, making this a very tiny object
-            Variables = null;
+            if (Variables != null) {
+                foreach (var value in Variables.Values)
+                    value.DecrementDreamObjectRefCount();
+
+                Variables = null;
+            }
 
             ObjectDefinition = null!;
         }
@@ -110,6 +121,8 @@ public void Delete() {
             if (TryGetProc("Del", out var delProc))
                 DreamThread.Run(delProc, this, null);
 
+            RefCount = 0;
+
             HandleDeletion();
         }
 
@@ -430,4 +443,24 @@ public override string ToString() {
             return ObjectDefinition.Type;
         }
     }
+
+    public static class DreamObjectExtensions {
+        [MethodImpl(MethodImplOptions.AggressiveInlining)]
+        public static void IncrementRefCount(this DreamObject dreamObject) {
+            Debug.Assert(!dreamObject.Deleted, "Invalid attempt at incrementing RefCount on a deleted DreamObject");
+
+            ++dreamObject.RefCount;
+        }
+
+        [MethodImpl(MethodImplOptions.AggressiveInlining)]
+        public static void DecrementRefCount(this DreamObject dreamObject) {
+            Debug.Assert(!dreamObject.Deleted, "Invalid attempt at decrementing RefCount on a deleted DreamObject");
+
+            var result = --dreamObject.RefCount;
+
+            Debug.Assert(result >= 0, "Invalid attempt at decrementing DreamObject's refcount when it is at (or below) 0");
+            if (result == 0)
+                dreamObject.Delete();
+        }
+    }
 }
diff --git a/OpenDreamRuntime/Objects/Types/DreamList.cs b/OpenDreamRuntime/Objects/Types/DreamList.cs
@@ -28,6 +28,13 @@ public DreamList(DreamObjectDefinition listDef, int size) : base(listDef) {
     private DreamList(DreamObjectDefinition listDef, List<DreamValue> values, Dictionary<DreamValue, DreamValue>? associativeValues) : base(listDef) {
         _values = values;
         _associativeValues = associativeValues;
+
+        IEnumerable<DreamValue> allContainedValues = values;
+        if (associativeValues != null)
+            allContainedValues = allContainedValues.Concat(associativeValues.Values);
+
+        foreach(var value in allContainedValues)
+            value.IncrementDreamObjectRefCount();
     }
 
     public override void Initialize(DreamProcArguments args) {
@@ -114,17 +121,26 @@ public virtual DreamValue GetValue(DreamValue key) {
     }
 
     public virtual void SetValue(DreamValue key, DreamValue value, bool allowGrowth = false) {
+        value.IncrementDreamObjectRefCount();
         if (key.TryGetValueAsInteger(out int keyInteger)) {
             if (allowGrowth && keyInteger == _values.Count + 1) {
                 _values.Add(value);
             } else {
                 _values[keyInteger - 1] = value;
             }
         } else {
+            key.IncrementDreamObjectRefCount();
             if (!ContainsValue(key)) _values.Add(key);
 
-            _associativeValues ??= new Dictionary<DreamValue, DreamValue>(1);
+            DreamValue? oldValueNullable = null;
+            if (_associativeValues == null)
+                _associativeValues = new Dictionary<DreamValue, DreamValue>(1);
+            else if (_associativeValues.TryGetValue(key, out var oldValue))
+                oldValueNullable = oldValue;
+
             _associativeValues[key] = value;
+            if (oldValueNullable.HasValue)
+                oldValueNullable.Value.DecrementDreamObjectRefCount();
         }
     }
 
@@ -138,6 +154,7 @@ public virtual void RemoveValue(DreamValue value) {
     }
 
     public virtual void AddValue(DreamValue value) {
+        value.IncrementDreamObjectRefCount();
         _values.Add(value);
     }
 
@@ -169,14 +186,19 @@ public virtual void Cut(int start = 1, int end = 0) {
         if (end == 0 || end > (_values.Count + 1)) end = _values.Count + 1;
 
         if (_associativeValues != null) {
-            for (int i = start; i < end; i++)
-                _associativeValues.Remove(_values[i - 1]);
+            for (int i = start; i < end; i++) {
+                var removedValue = _values[i - 1];
+                if (_associativeValues.Remove(removedValue, out var removedAssoc))
+                    removedAssoc.DecrementDreamObjectRefCount();
+                removedValue.DecrementDreamObjectRefCount();
+            }
         }
 
         _values.RemoveRange(start - 1, end - start);
     }
 
     public void Insert(int index, DreamValue value) {
+        value.IncrementDreamObjectRefCount();
         _values.Insert(index - 1, value);
     }
 
@@ -375,6 +397,18 @@ public override DreamValue OperatorEquivalent(DreamValue b) {
 
         return DreamValue.True;
     }
+
+    protected override void HandleDeletion() {
+        base.HandleDeletion();
+
+        if (_associativeValues != null)
+            foreach (var assocValue in _associativeValues.Values)
+                assocValue.DecrementDreamObjectRefCount();
+
+        foreach (var value in _values)
+            value.DecrementDreamObjectRefCount();
+    }
+
     #endregion Operators
 }
 
@@ -442,6 +476,12 @@ public override DreamValue Initial(string name) {
     public override bool IsSaved(string name) {
         return _dreamObject.IsSaved(name);
     }
+
+    protected override void HandleDeletion() {
+        base.HandleDeletion();
+
+        _dreamObject.DecrementRefCount();
+    }
 }
 
 // global.vars list

diff --git a/OpenDreamRuntime/Procs/AsyncNativeProc.cs b/OpenDreamRuntime/Procs/AsyncNativeProc.cs
@@ -49,6 +49,13 @@ public void Initialize(AsyncNativeProc? proc, Func<State, Task<DreamValue>> task
                 Usr = usr;
                 arguments.Values.CopyTo(_arguments);
                 _argumentCount = arguments.Count;
+
+                src?.IncrementRefCount();
+
+                usr?.IncrementRefCount();
+
+                foreach (var arg in arguments.Values)
+                    arg.IncrementDreamObjectRefCount();
             }
 
             // Used to avoid reentrant resumptions in our proc
@@ -169,7 +176,7 @@ public DreamValue GetArgument(int argumentPosition, string argumentName) {
         private readonly Dictionary<string, DreamValue>? _defaultArgumentValues;
         private readonly Func<State, Task<DreamValue>> _taskFunc;
 
-        public AsyncNativeProc(int id, TreeEntry owningType, string name, List<string> argumentNames, Dictionary<string, DreamValue> defaultArgumentValues, Func<State, Task<DreamValue>> taskFunc)
+        public AsyncNativeProc(int id, TreeEntry owningType, string name, List<string> argumentNames, Dictionary<string, DreamValue>? defaultArgumentValues, Func<State, Task<DreamValue>> taskFunc)
             : base(id, owningType, name, null, ProcAttributes.None, argumentNames, null, null, null, null, 0) {
             _defaultArgumentValues = defaultArgumentValues;
             _taskFunc = taskFunc;

diff --git a/OpenDreamRuntime/Procs/DMProc.cs b/OpenDreamRuntime/Procs/DMProc.cs
@@ -361,6 +361,8 @@ private DMProcState(DMProcState other, DreamThread thread) {
             _stack = _dreamValuePool.Rent(other._stack.Length);
             _localVariables = _dreamValuePool.Rent(other._localVariables.Length);
             Array.Copy(other._localVariables, _localVariables, other._localVariables.Length);
+
+            IncrementRefCounts();
         }
 
         public void Initialize(DMProc proc, DreamThread thread, int maxStackSize, DreamObject? instance, DreamObject? usr, DreamProcArguments arguments) {
@@ -376,6 +378,24 @@ public void Initialize(DMProc proc, DreamThread thread, int maxStackSize, DreamO
             for (int i = 0; i < ArgumentCount; i++) {
                 _localVariables[i] = arguments.GetArgument(i);
             }
+
+            IncrementRefCounts();
+        }
+
+        private void IncrementRefCounts() {
+            Instance?.IncrementRefCount();
+            Usr?.IncrementRefCount();
+
+            for(var i = 0; i < ArgumentCount; ++i)
+                _localVariables[i].IncrementDreamObjectRefCount();
+        }
+
+        private void DecrementRefCounts() {
+            for(var i = ArgumentCount - 1; i >= 0; --i)
+                _localVariables[i].DecrementDreamObjectRefCount();
+
+            Usr?.DecrementRefCount();
+            Instance?.DecrementRefCount();
         }
 
         public override unsafe ProcStatus Resume() {
@@ -421,6 +441,7 @@ public override unsafe ProcStatus Resume() {
         }
 
         public override void ReturnedInto(DreamValue value) {
+            value.IncrementDreamObjectRefCount();
             Push(value);
         }
 
@@ -442,6 +463,8 @@ public void Jump(int position) {
         }
 
         public void SetReturn(DreamValue value) {
+            value.IncrementDreamObjectRefCount();
+            Result.DecrementDreamObjectRefCount();
             Result = value;
         }
 
@@ -510,6 +533,7 @@ public override void Dispose() {
             _stackIndex = 0;
             _stack = null;
 
+            DecrementRefCounts();
             _dreamValuePool.Return(_localVariables);
             _localVariables = null;
 
@@ -527,6 +551,8 @@ public void SetArgument(int id, DreamValue value) {
             if (id < 0 || id >= ArgumentCount)
                 throw new IndexOutOfRangeException($"Given argument id ({id}) was out of range");
 
+            value.IncrementDreamObjectRefCount();
+            _localVariables[id].DecrementDreamObjectRefCount();
             _localVariables[id] = value;
         }
 
@@ -699,24 +725,49 @@ private static void ThrowReferenceNotListIndex() {
         }
 
         public void AssignReference(DreamReference reference, DreamValue value) {
+            DreamObject? dreamObject;
             switch (reference.Type) {
-                case DMReference.Type.Self: Result = value; break;
+                case DMReference.Type.Self: Result = value;
+                    value.IncrementDreamObjectRefCount();
+                    Result.DecrementDreamObjectRefCount();
+                    Result = value;
+                    break;
                 case DMReference.Type.Argument: SetArgument(reference.Value, value); break;
-                case DMReference.Type.Local: _localVariables[ArgumentCount + reference.Value] = value; break;
+                case DMReference.Type.Local:
+                    var localIndex = ArgumentCount + reference.Value;
+                    value.IncrementDreamObjectRefCount();
+                    // have to guard against rented previous state here
+                    if (_localVariables[localIndex].TryGetValueAsDreamObject(out dreamObject)
+                        && dreamObject?.Deleted == false)
+                        dreamObject.DecrementRefCount();
+
+                    _localVariables[localIndex] = value;
+                    break;
                 case DMReference.Type.SrcField: Instance.SetVariable(ResolveString(reference.Value), value); break;
-                case DMReference.Type.Global: DreamManager.Globals[reference.Value] = value; break;
+                case DMReference.Type.Global:
+                    value.IncrementDreamObjectRefCount();
+                    DreamManager.Globals[reference.Value].DecrementDreamObjectRefCount();
+                    DreamManager.Globals[reference.Value] = value;
+                    break;
                 case DMReference.Type.Src:
+                    value.IncrementDreamObjectRefCount();
                     //TODO: src can be assigned to non-DreamObject values
+                    var oldInstance = Instance;
                     if (!value.TryGetValueAsDreamObject(out Instance)) {
                         ThrowCannotAssignSrcTo(value);
                     }
 
+                    oldInstance?.DecrementRefCount();
                     break;
                 case DMReference.Type.Usr:
+                    value.IncrementDreamObjectRefCount();
                     //TODO: usr can be assigned to non-DreamObject values
+                    var oldUsr = Usr;
                     if (!value.TryGetValueAsDreamObject(out Usr)) {
                         ThrowCannotAssignUsrTo(value);
                     }
+
+                    oldUsr?.DecrementRefCount();
                     break;
                 case DMReference.Type.Field: {
                     DreamValue owner = Pop();
@@ -729,7 +780,7 @@ public void AssignReference(DreamReference reference, DreamValue value) {
                 case DMReference.Type.ListIndex: {
                     GetIndexReferenceValues(reference, out var index, out var indexing);
 
-                    if (indexing.TryGetValueAsDreamObject(out var dreamObject) && dreamObject != null) {
+                    if (indexing.TryGetValueAsDreamObject(out dreamObject) && dreamObject != null) {
                         dreamObject.OperatorIndexAssign(index, value);
                     } else {
                         ThrowCannotAssignListIndex(index, indexing);

diff --git a/OpenDreamRuntime/Procs/InitDreamObject.cs b/OpenDreamRuntime/Procs/InitDreamObject.cs
@@ -32,6 +32,11 @@ public void Initialize(DreamThread thread, DreamObject dreamObject, DreamObject?
             arguments.Values.CopyTo(_arguments);
             _argumentCount = arguments.Count;
             _stage = Stage.Init;
+
+            dreamObject.IncrementRefCount();
+            usr?.IncrementRefCount();
+            foreach(var arg in arguments.Values)
+                arg.IncrementDreamObjectRefCount();
         }
 
         private DreamObject _dreamObject;

diff --git a/OpenDreamRuntime/Procs/Native/DreamProcNative.cs b/OpenDreamRuntime/Procs/Native/DreamProcNative.cs
@@ -70,6 +70,7 @@ public static void SetupNativeProcs(DreamObjectTree objectTree) {
             objectTree.SetGlobalNativeProc(DreamProcNativeRoot.NativeProc_rand_seed);
             objectTree.SetGlobalNativeProc(DreamProcNativeRoot.NativeProc_range);
             objectTree.SetGlobalNativeProc(DreamProcNativeRoot.NativeProc_ref);
+            objectTree.SetGlobalNativeProc(DreamProcNativeRoot.NativeProc_refcount);
             objectTree.SetGlobalNativeProc(DreamProcNativeRoot.NativeProc_regex);
             objectTree.SetGlobalNativeProc(DreamProcNativeRoot.NativeProc_replacetext);
             objectTree.SetGlobalNativeProc(DreamProcNativeRoot.NativeProc_replacetextEx);