Make CheckStateCount() behave as expected (#210)

src/Runtime/Distributions/Automata/Automaton.cs

@@ -2634,14 +2634,14 @@ namespace Microsoft.ML.Probabilistic.Distributions.Automata
         /// </summary>
         public class UnlimitedStatesComputation : IDisposable
         {
-            private readonly int originalMaxStateCount;
+            private readonly int originalThreadMaxStateCount;
 
             /// <summary>
             /// Initializes a new instance of the <see cref="UnlimitedStatesComputation"/> class.
             /// </summary>
             public UnlimitedStatesComputation()
             {
-                originalMaxStateCount = threadMaxStateCountOverride;
+                this.originalThreadMaxStateCount = threadMaxStateCountOverride;
                 threadMaxStateCountOverride = int.MaxValue;
             }
 
@@ -2650,15 +2650,18 @@ namespace Microsoft.ML.Probabilistic.Distributions.Automata
             /// </summary>
             public void CheckStateCount(TThis automaton)
             {
-                if (automaton.States.Count > originalMaxStateCount)
+                var limit = this.originalThreadMaxStateCount != 0
+                    ? this.originalThreadMaxStateCount
+                    : maxStateCount;
+                if (automaton.States.Count > limit)
                 {
-                    throw new AutomatonTooLargeException(originalMaxStateCount);
+                    throw new AutomatonTooLargeException(limit);
                 }
             }
 
             public void Dispose()
             {
-                threadMaxStateCountOverride = originalMaxStateCount;
+                threadMaxStateCountOverride = this.originalThreadMaxStateCount;
             }
         }
         #endregion

test/Tests/Strings/AutomatonTests.cs

@@ -884,6 +884,40 @@ namespace Microsoft.ML.Probabilistic.Tests
             }
         }
 
+        /// <summary>
+        /// Tests whether StringAutomaton.UnlimitedStatesComputation.CheckStateCount() works as expected
+        /// </summary>
+        [Fact]
+        [Trait("Category", "StringInference")]
+        public void CheckStateCount()
+        {
+            using (var unlimited = new StringAutomaton.UnlimitedStatesComputation())
+            {
+                var builder = new StringAutomaton.Builder();
+                var state = builder.Start;
+
+                for (var i = 1; i < 200000; ++i)
+                {
+                    state = state.AddTransition('a', Weight.One);
+                }
+
+                var automaton = builder.GetAutomaton();
+
+                // Fine, because 200k < default limit
+                unlimited.CheckStateCount(automaton);
+
+                for (var i = 1; i < 200000; ++i)
+                {
+                    state = state.AddTransition('a', Weight.One);
+                }
+
+                automaton = builder.GetAutomaton();
+
+                // Not fine anymore, automaton (with 400k states) is over the default limit
+                Assert.Throws<AutomatonTooLargeException>(() => unlimited.CheckStateCount(automaton));
+            }
+        }
+
         /// <summary>
         /// Tests creating an automaton from state and transition lists.
         /// </summary>