disposed state of DisposeScope & linked list like DisposeScopeStack

Author: yueyinqiu

src/TorchSharp/DisposeScope.cs

@@ -14,39 +14,47 @@ namespace TorchSharp
     /// </summary>
     public sealed class DisposeScope : IDisposable
     {
-        private readonly DisposeScopeManager _disposeScopeManager;
+        private DisposeScopeManager? _disposeScopeManager;
 
-        public DisposeScope(DisposeScopeManager disposeScopeManager)
+        internal DisposeScope(DisposeScopeManager disposeScopeManager)
         {
-            _disposeScopeManager = disposeScopeManager ?? throw new ArgumentNullException(nameof(disposeScopeManager));
-            if (disposeScopeManager.DisposeScopeStack.Count > 0) {
-                OuterScope = disposeScopeManager.DisposeScopeStack[
-                    disposeScopeManager.DisposeScopeStack.Count - 1];
-            }
+            _disposeScopeManager = disposeScopeManager;
+            this.OuterScope = disposeScopeManager.CurrentDisposeScope;
         }
 
         /// <summary>
         /// The outer scope with relation to this scope.
         /// </summary>
-        internal DisposeScope? OuterScope { get; }
+        internal DisposeScope? OuterScope { get; set; }
 
         /// <summary>
         /// The disposables that are scheduled for disposing.
         /// </summary>
-        /// TODO: There is a ReferenceEqualityComparer coming in .NET 6, use that!
         internal HashSet<IDisposable> Disposables { get; private set; } =
             new HashSet<IDisposable>(ReferenceEqualityComparer<IDisposable>.Default);
 
         /// <summary>
         /// A view of the disposables in the scope - this list will not be kept in synch with the disposables
         /// in the scope.
         /// </summary>
-        public IReadOnlyList<IDisposable> DisposablesView => Disposables.ToList();
+        public IReadOnlyList<IDisposable> DisposablesView {
+            get {
+                if (this._disposeScopeManager is null)
+                    throw new ObjectDisposedException("The dispose scope has been disposed.");
+                return Disposables.ToArray();
+            }
+        }
 
         /// <summary>
         /// The number of disposables currently held in the scope
         /// </summary>
-        public int DisposablesCount => Disposables.Count;
+        public int DisposablesCount {
+            get {
+                if (this._disposeScopeManager is null)
+                    throw new ObjectDisposedException("The dispose scope has been disposed.");
+                return Disposables.Count;
+            }
+        }
 
         /// <summary>
         /// Includes a disposable in the scope - for tensors this is done automatically once the scope has been
@@ -57,6 +65,8 @@ public DisposeScope(DisposeScopeManager disposeScopeManager)
         /// <returns></returns>
         public T Include<T>(T disposable) where T : IDisposable
         {
+            if (this._disposeScopeManager is null)
+                throw new ObjectDisposedException("The dispose scope has been disposed.");
             Disposables.Add(disposable);
             return disposable;
         }
@@ -157,6 +167,8 @@ public void MoveToOther(DisposeScope? scope, params IDisposable[] disposables) =
         /// </summary>
         public void MoveToOther(DisposeScope? scope, IEnumerable<IDisposable> disposables)
         {
+            if (this._disposeScopeManager is null)
+                throw new ObjectDisposedException("The dispose scope has been disposed.");
             foreach (var disposable in disposables) {
                 if (Disposables.Remove(disposable)) {
                     AddToOther(scope, disposable);
@@ -208,6 +220,8 @@ public T Detach<T>(T disposable) where T : IDisposable
         /// </summary>
         public void Detach(IEnumerable<IDisposable> disposables)
         {
+            if (this._disposeScopeManager is null)
+                throw new ObjectDisposedException("The dispose scope has been disposed.");
             foreach (var disposable in disposables) {
                 if (Disposables.Remove(disposable)) {
                     _disposeScopeManager.StatisticsInstance.DetachedFromScopeCount++;
@@ -220,6 +234,8 @@ public void Detach(IEnumerable<IDisposable> disposables)
 
         public void Attach(IDisposable disposable)
         {
+            if (this._disposeScopeManager is null)
+                throw new ObjectDisposedException("The dispose scope has been disposed.");
             if (disposable is torch.Tensor tensor) {
                 if (tensor.OwningDisposeScope == null && !tensor.IsInvalid) {
                     _disposeScopeManager.StatisticsInstance.DetachedFromScopeCount--;
@@ -241,6 +257,8 @@ public void Attach(IDisposable disposable)
         /// </summary>
         public void DisposeEverythingBut(IEnumerable<IDisposable> inKeep)
         {
+            if (this._disposeScopeManager is null)
+                throw new ObjectDisposedException("The dispose scope has been disposed.");
             // Avoiding multiple enumerations
             var oldList = Disposables;
             Disposables = inKeep.ToHashSet(ReferenceEqualityComparer<IDisposable>.Default);
@@ -316,8 +334,11 @@ public T DisposeEverythingBut<T>(T keep) where T : IDisposable
         /// </summary>
         public void Dispose()
         {
+            if (this._disposeScopeManager is null)
+                return;
             DisposeEverything();
             _disposeScopeManager.RemoveDisposeScope(this);
+            this._disposeScopeManager = null;
         }
 
         /// <summary>
@@ -329,6 +350,8 @@ public void Dispose()
         /// <param name="disposable">The disposable that was disposed</param>
         public void MarkAsDisposed(IDisposable disposable)
         {
+            if (this._disposeScopeManager is null)
+                throw new ObjectDisposedException("The dispose scope has been disposed.");
             _disposeScopeManager.StatisticsInstance.DisposedInScopeCount++;
             Disposables.Remove(disposable);
             if (disposable is torch.Tensor tensor) {
@@ -345,6 +368,8 @@ public void MarkAsDisposed(IDisposable disposable)
 
         private void AddToOther(DisposeScope? scope, IDisposable disposable)
         {
+            if (this._disposeScopeManager is null)
+                throw new ObjectDisposedException("The dispose scope has been disposed.");
             if (scope != null) {
                 scope.Disposables.Add(disposable);
             } else {

src/TorchSharp/DisposeScopeManager.cs

@@ -18,21 +18,20 @@ public class DisposeScopeManager
         internal ThreadDisposeScopeStatistics StatisticsInstance { get; } = new ThreadDisposeScopeStatistics();
 
         internal static DisposeScopeManager ThreadSingleton => (_threadSingleton ??= new DisposeScopeManager());
-        internal List<DisposeScope> DisposeScopeStack { get; } = new();
+        internal DisposeScope? CurrentDisposeScope { get; set; } = null;
 
         public static ThreadDisposeScopeStatistics Statistics => ThreadSingleton.StatisticsInstance;
 
         internal DisposeScope? RegisterOnCurrentDisposeScope(IDisposable disposable)
         {
-            if (DisposeScopeStack.Count == 0) {
+            if (this.CurrentDisposeScope is null) {
                 StatisticsInstance.CreatedOutsideScopeCount++;
                 return null;
             }
 
             StatisticsInstance.CreatedInScopeCount++;
-            var current = DisposeScopeStack[DisposeScopeStack.Count - 1];
-            current.Include(disposable);
-            return current;
+            this.CurrentDisposeScope.Include(disposable);
+            return CurrentDisposeScope;
         }
 
         internal static DisposeScope NewDisposeScope()
@@ -42,16 +41,33 @@ internal static DisposeScope NewDisposeScope()
 
         internal void RemoveDisposeScope(DisposeScope disposeScope)
         {
-            var index = DisposeScopeStack.LastIndexOf(disposeScope);
-            if (index is not -1)
-                DisposeScopeStack.RemoveAt(index);
+            var scope = this.CurrentDisposeScope;
+            if (object.ReferenceEquals(scope, disposeScope)) {
+                this.CurrentDisposeScope = scope.OuterScope;
+                return;
+            }
+            if (scope is null) {
+                return;
+            }
+
+            for (; ; ) {
+                var outerScope = scope.OuterScope;
+                if (object.ReferenceEquals(outerScope, disposeScope)) {
+                    scope.OuterScope = outerScope.OuterScope;
+                    return;
+                }
+
+                if (outerScope is null) {
+                    return;
+                }
+                scope = outerScope;
+            }
         }
 
         private DisposeScope InnerNewDisposeScope()
         {
-            var disposeScope = new DisposeScope(this);
-            DisposeScopeStack.Add(disposeScope);
-            return disposeScope;
+            this.CurrentDisposeScope = new DisposeScope(this);
+            return this.CurrentDisposeScope;
         }
     }
 }