Resolve parameteric sorts in ranks, and ArraysEx functions

Author: kjcjohnson

Semgus-Lib/Model/Smt/SmtFunction.cs

@@ -153,7 +153,22 @@ private bool TryResolveRank([NotNullWhen(true)] out SmtFunctionRank? rank, SmtFu
                     }
                     else
                     {
-                        resolvedParameters.Add(templateSort, concreteSort);
+                        // Template sorts (that are a parameter) are allowed to have arity 0.
+                        // The whole template will just be matched with the concrete sort.
+                        // Otherwise, we want to make sure the arities match.
+                        if (!(templateSort.IsSortParameter && templateSort.Arity == 0)
+                            && templateSort.Arity != concreteSort.Arity)
+                        {
+                            rank = default;
+                            return false;
+                        }
+
+                        // Recursively match the template and match parameter sorts
+                        if (!TraverseAndMatchTemplate(templateSort, concreteSort, resolvedParameters))
+                        {
+                            rank = default;
+                            return false;
+                        }
                         templateSort = concreteSort;
                     }
                 }
@@ -212,6 +227,62 @@ private bool TryResolveRank([NotNullWhen(true)] out SmtFunctionRank? rank, SmtFu
             return template.Validator(rank);
         }
 
+        /// <summary>
+        /// Match a template sort and parameters against a concrete sort
+        /// </summary>
+        /// <param name="template">Template sort to match</param>
+        /// <param name="concrete">Concrete sort to match against</param>
+        /// <param name="resolvedParameters">Dictionary of resolved parameters</param>
+        /// <returns>True if successfully resolves, false if not</returns>
+        private static bool TraverseAndMatchTemplate(SmtSort template, SmtSort concrete, IDictionary<SmtSort, SmtSort> resolvedParameters)
+        {
+            // Template can be a parameter, and concrete can be whatever
+            if (template.IsSortParameter)
+            {
+                if (resolvedParameters.TryGetValue(template, out var resolved))
+                {
+                    if (resolved != concrete)
+                    {
+                        return false;
+                    }
+                }
+                else
+                {
+                    resolvedParameters.Add(template, concrete);
+                }
+
+                if (template.Arity == 0 && concrete.Arity > 0)
+                {
+                    return true;
+                }
+            }
+
+            // Otherwise, arities must match, so we can traverse
+            if (template.Arity != concrete.Arity)
+            {
+                return false;
+            }
+
+            if (template.IsSortParameter)
+            {
+                foreach (var (tempParam, concParam) in template.Parameters.Zip(concrete.Parameters))
+                {
+                    if (!TraverseAndMatchTemplate(tempParam, concParam, resolvedParameters))
+                    {
+                        return false;
+                    }
+                }
+            }
+            else
+            {
+                if (template != concrete)
+                {
+                    return false;
+                }
+            }
+            return true;
+        }
+
         /// <summary>
         /// Returns a string describing available ranks
         /// </summary>

Semgus-Lib/Model/Smt/SmtSort.cs

@@ -47,6 +47,11 @@ public SmtSort(SmtSortIdentifier name)
         /// <param name="resolved">Resolved parameters. Should have same length as arity</param>
         public virtual void UpdateForResolvedParameters(IList<SmtSort> resolved) { }
 
+        /// <summary>
+        /// This sort's parameters
+        /// </summary>
+        public virtual IEnumerable<SmtSort> Parameters { get; } = Enumerable.Empty<SmtSort>();
+
         /// <summary>
         /// An arbitrary generic sort
         /// </summary>

Semgus-Lib/Model/Smt/Theories/SmtArraysExTheory.cs

@@ -42,6 +42,14 @@ internal sealed class ArraySort : SmtSort
             /// </summary>
             public SmtSort ValueSort { get; private set; }
 
+            /// <summary>
+            /// This sort's parameters
+            /// </summary>
+            public override IEnumerable<SmtSort> Parameters
+            {
+                get => new SmtSort[] { IndexSort, ValueSort };
+            }
+
             /// <summary>
             /// Constructs a new array sort with the given parameters
             /// </summary>
@@ -55,6 +63,20 @@ private ArraySort(SmtSortIdentifier indexSort, SmtSortIdentifier valueSort) :
                 Arity = 2;
             }
 
+            /// <summary>
+            /// Constructs an array sort with specific sort objects
+            /// </summary>
+            /// <param name="indexSort">The sort for the index</param>
+            /// <param name="valueSort">The sort for the value</param>
+            internal ArraySort(SmtSort indexSort, SmtSort valueSort)
+                : base(new(ArraySortPrimaryId, indexSort.Name, valueSort.Name))
+            {
+                IndexSort = indexSort;
+                ValueSort = valueSort;
+                IsSortParameter = indexSort.IsSortParameter || valueSort.IsSortParameter;
+                Arity = 2;
+            }
+
             /// <summary>
             /// Gets the array sort for the given index and value sorts
             /// </summary>
@@ -117,8 +139,12 @@ public override void UpdateForResolvedParameters(IList<SmtSort> resolved)
         private SmtArraysExTheory()
         {
             SmtSourceBuilder sb = new(this);
-            sb.AddOnTheFlyFn("select");
-            sb.AddOnTheFlyFn("store");
+            var usf = new SmtSort.UniqueSortFactory();
+            var valueSort = usf.Next();
+            var indexSort = usf.Next();
+            var arraySort = new ArraySort(indexSort, valueSort);
+            sb.AddFn("select", valueSort, arraySort, indexSort);
+            sb.AddFn("store", arraySort, arraySort, indexSort, valueSort);
 
             Functions = sb.Functions;
             PrimaryFunctionSymbols = sb.PrimaryFunctionSymbols;