Cat - A Statically Typed Programming Language Interpreter in C#

/// Public domain code by Christopher Diggins
/// http://www.cat-language.com

using System;
using System.Collections.Generic;
using System.Text;

namespace Cat
{
    /// <summary>
    /// Type variable declarations and type variables which are the result of construction by 
    /// inference might have constraints applied to them. For example, in the function [dup +] 
    /// The type of dup is (A:any*)->(A A) but the addition operation implies a constraint that
    /// A must be an integer. The constraintSolver generates constraints through a process called
    /// integration. You then solving the set of equations using MergeAllConstraint.  
    /// Finally you use ResolveFunction to replace type variables in the two equations with 
    /// the results which can be concrete types or a reduced set of variables.
    /// </summary>
    public class ConstraintSolver : CatBase
    {
        class Constraints : Dictionary<TypeVarDecl, List<CatType>>
        {
            /// <summary>
            /// Overridden for debugging purposes. 
            /// </summary>
            public override string ToString()
            {
                string ret = "";
                foreach (TypeVarDecl decl in Keys)
                {
                    List<CatType> list = this[decl];
                    ret += "(" + decl.ToString() + " = ";
                    foreach (CatType t in list)
                        ret += t.ToString() + ";";                       
                    ret += ") ";
                                        
                }
                return ret;
            }
        }

        /// <summary>
        /// The ConstraintIterator is a recursive iterator which is aware of constraints
        /// It simplifies the logic of integrating constraints. 
        /// </summary>
        class ConstraintIterator : CatBase
        {
            TypeIterator mIter;
            Stack<TypeIterator> mIterStack = new Stack<TypeIterator>();
            ConstraintSolver mSolver;

            public ConstraintIterator(TypeIterator iter, ConstraintSolver solver)
            {
                mIter = iter;
                mSolver = solver;
                GoDeep();
            }

            public void GotoNext()
            {
                if (mIter.AtEnd())
                {
                    GoShallow();
                }
                else
                {
                    mIter.GotoNext();
                    GoDeep();
                }
            }

            private void GoDeep()
            {
                if (!AtEnd())
                {
                    CatType t = GetValue();
                    if (mSolver.HasStackConstraint(t))
                    {
                        TypeStack stk = mSolver.GetStackConstraint(t);
                        mIterStack.Push(mIter);
                        mIter = stk.GetTypeIterator();
                        GoDeep();
                    }
                }
                else
                {
                    GoShallow();
                }
            }

            private void GoShallow()
            {
                Assert(AtEnd());
                if (mIterStack.Count > 0)
                {
                    mIter = mIterStack.Pop();
                    GotoNext();
                }
            }

            public CatType GetValue()
            {
                Assert(!AtEnd());
                CatType t = mIter.GetValue();
                return t;    
            }

            public bool AtEnd()
            {
                return mIter.AtEnd();
            }

            public override string ToString()
            {
                return mIter.ToString();
            }
        }

        #region fields 
        Constraints mConstraints = new Constraints();
        bool mbConstraintsMerged = false;
        #endregion

        #region private functions
        CatType GetConstraint(CatType t)
        {
            Assert(t.IsTypeVar(), "expected type variable, for call to ResolveTypeVar");
            if (!HasConstraint(t.AsDecl()))
                return t;
            List<CatType> list = mConstraints[t.AsDecl()];
            if (list.Count == 0)
                return t;
            Assert(list.Count == 1, "found multiple constrainted for type variable " + t.ToString() + ", perhaps missing a call to merge constraints");
            return list[0];
        }
        void AddConstraint(CatType src, CatType constraint)
        {
            Assert(src.IsTypeVar());
            Assert(!(constraint is TypeStack));
            List<CatType> list = null;
            if (!HasConstraint(src.AsDecl()))
            {
                list = new List<CatType>();
                mConstraints.Add(src.AsDecl(), list);
            }
            else
            {
                // If there are any variables ... replacing them if this is a value?
                // or does it matter?
                list = mConstraints[src.AsDecl()];
            }
            list.Add(constraint);
        }
        void AddConstraint(CatType src, TypeStack stack)
        {
            Assert(src.IsTypeVar());
            Assert(src.IsMultiVar());
            List<CatType> list = null;
            if (!HasConstraint(src.AsDecl()))
            {
                list = new List<CatType>();
                mConstraints.Add(src.AsDecl(), list);
            }
            else
            {
                list = mConstraints[src.AsDecl()];
            }
            list.Add(stack);
        }
        TypeStack GetStackConstraint(CatType src)
        {
            Assert(HasStackConstraint(src));
            List<CatType> list = mConstraints[src.AsDecl()];
            // NOTE: I am not happy about this!
            Assert(list.Count >= 1);
            CatType ret = list[0];
            Assert(ret is TypeStack);
            return ret as TypeStack;
        }
        bool HasConstraint(CatType src)
        {
            return mConstraints.ContainsKey(src.AsDecl());
        }
        /// <summary>
        /// Returns any concrete constraint associated with a particular variable.
        /// This has a side-effect that the list of constraints for the variable 
        /// is shortened if a type is found. This is to reduce the overall complexity
        /// of the constraint merging algorithm.
        /// </summary>
        CatType GetConcreteConstraint(CatType src)
        {
            if (!mConstraints.ContainsKey(src.AsDecl()))
                return null;
            List<CatType> list = mConstraints[src.AsDecl()];

            CatType ret = null;
            foreach (CatType t in list) 
            {
                if (!t.IsTypeVar())
                    ret = GetBetterConstraint(ret, t);
            }

            // If there was a type found, make sure that the list of constraints only 
            // contains it.
            if (ret != null && list.Count > 1)
            {                
                list.Clear();
                list.Add(ret);
            }

            return ret;
        }
        bool HasStackConstraint(CatType src)
        {
            return src.IsMultiVar() && HasConstraint(src);
        }
        void SetConstraint(CatType src, CatType constraint)
        {
            Assert(src.IsSingleVar());
            AddConstraint(src, constraint);
        }
        void SetConstraint(CatType src, ConstraintIterator iter)
        {   
            // TODO: add flags for when this stuff is in a sub-function ... maybe..
            Assert(src.IsMultiVar(), "expected a multi-variable type");

            TypeStack stk = MakeStackFromIter(iter);
            AddConstraint(src, stk);
        }
        /// <summary>
        /// Construct a TypeStack using an iterator, starting at its current position. 
        /// This function changes the iterator state, so that at the end it will always
        /// have AtEnd() == true
        /// </summary>
        TypeStack MakeStackFromIter(ConstraintIterator iter)
        {
            TypeStack ret = new TypeStack();
            for (; !iter.AtEnd(); iter.GotoNext())
                ret.PushFront(iter.GetValue());
            return ret;
        }
        void IntegrateFxnConstraints(FxnType ftLeft, FxnType ftRight)
        {
            IntegrateConstraints(ftLeft.GetMainConsIter(), ftRight.GetMainConsIter());
            IntegrateConstraints(ftLeft.GetAuxConsIter(), ftRight.GetAuxConsIter());
            IntegrateConstraints(ftLeft.GetMainProdIter(), ftRight.GetMainProdIter());
            IntegrateConstraints(ftLeft.GetAuxProdIter(), ftRight.GetAuxProdIter());
        }
        private CatType Resolve(CatType t)
        {
            Assert(mbConstraintsMerged, "missing call to MergeAllConstraints");
            if (t is FxnType)
            {
                FxnType f = t as FxnType;
                Resolve(f.GetMainConsStack());
                Resolve(f.GetAuxConsStack());
                Resolve(f.GetMainProdStack());
                Resolve(f.GetAuxProdStack());
            }
            else if (t.IsTypeVar())
            {
                return GetConstraint(t);
            }
            else if (t is TypeStack)
            {
                TypeStack stk = t as TypeStack;
                for (int i = stk.Count - 1; i >= 0; --i)
                    stk.ReplaceAndExpand(i, Resolve(stk[i]));
            }
            else if (t is SimpleType)
            {
                return t;
            }
            else
            {
                Assert(false, "unhandled type " + t.ToString() + " in call to InternalResolve");
            }
            return t;
        }
        #endregion

        #region public functions
        /// <summary>
        /// This will identify and store any variable constraints based on relative position
        /// in the two stacks. 
        /// </summary>
        public void IntegrateConstraints(TypeIterator iterLeft, TypeIterator iterRight)
        {
            IntegrateConstraints(new ConstraintIterator(iterLeft, this), new ConstraintIterator(iterRight, this));
        }
        private void IntegrateConstraints(ConstraintIterator iterLeft, ConstraintIterator iterRight)
        {
            Assert(!mbConstraintsMerged, "constraints are already merged");
            while (!iterLeft.AtEnd() || !iterRight.AtEnd())
            {
                CatType right = iterRight.AtEnd() ? null : iterRight.GetValue();
                CatType left = iterLeft.AtEnd() ? null : iterLeft.GetValue();

                if (right == null)
                {
                    if (left.IsMultiVar())
                    {
                        SetConstraint(left, iterRight);
                    }
                    else
                    {
                        throw new Exception("expected " + left.ToString() + " but encountered end of stack");
                    }
                }
                else if (left == null)
                {
                    if (right.IsMultiVar())
                    {
                        SetConstraint(right, iterLeft);
                    }
                    else
                    {
                        throw new Exception("unexpected end of stack");
                    }
                }
                else if (right.IsMultiVar() && right is TypeVarDecl)
                {
                    SetConstraint(right, iterLeft);
                }
                else if (left.IsMultiVar())
                {
                    SetConstraint(left, iterRight);
                }
                else if (right.IsSingleVar())
                {
                    SetConstraint(right, left);
                }
                else if (left.IsSingleVar())
                {
                    SetConstraint(left, right);
                }
                else if (right is FxnType)
                {
                    if (left is FxnType)
                    {
                        IntegrateFxnConstraints(left as FxnType, right as FxnType);
                    }
                    else
                    {
                        throw new Exception("type constraint error: left hand side must be a function or type variable");
                    }
                }
                else
                {
                    // We do very basic type checking, to make sure both types are the same.
                    // TODO: this should be improved.
                    if (!right.IsTypeVar() && !left.IsTypeVar())
                    {
                        if (!right.IsTypeEq(left))
                            Throw("mismatched types " + right.ToString() + " and " + left.ToString());
                    }
                }

                // Increment the iterators if safe to do so 
                if (!iterLeft.AtEnd())
                    iterLeft.GotoNext();
                if (!iterRight.AtEnd())
                    iterRight.GotoNext();
            }
        }

        private void ResolveConstraintVars(CatType t, List<CatType> dest, List<CatType> visited)
        {
            Assert(t.IsTypeVar(), "expected a TypeVar instead of " + t.ToString());
            if (visited.Contains(t))
                return;
            visited.Add(t);
            if (HasConstraint(t.AsDecl()))
            {
                List<CatType> src = mConstraints[t.AsDecl()];
                ResolveConstraintVars(src, dest, visited);
                dest.AddRange(src);
            }
        }

        private void ResolveConstraintVars(List<CatType> src, List<CatType> dest, List<CatType> visited)
        {
            for (int i=src.Count-1; i >= 0; --i)
            {
                CatType t = src[i];
                if (t.IsTypeVar())
                {
                    ResolveConstraintVars(t, dest, visited);
                }
                else
                {
                    dest.Add(t);
                }
            }
        }
        private void ResolveConstraintVars()
        {
            List<CatType> visited = new List<CatType>();
            foreach (CatType t in mConstraints.Keys)
            {
                List<CatType> src = mConstraints[t.AsDecl()];
                List<CatType> dest = new List<CatType>();
                ResolveConstraintVars(src, dest, visited);
                src.AddRange(dest);
            }
        }
        /// <summary>
        /// Returns true if x is a subtype of y. This implies that if x == y then
        /// this is also true.
        /// </summary>
        /// <param name="x">sub-type</param>
        /// <param name="y">super-type</param>
        /// <returns>true iff x is a subtype of y</returns>
        private bool IsSubType(CatType x, CatType y)
        {
            if (x == y)
            {
                return true;
            }
            else if (y.IsTypeVar())
            {
                return true;
            }
            else if (x.IsTypeVar())
            {
                return false;
            }
            else if (x is SimpleType)
            {
                Assert(x.IsTypeEq(y), "expected simple type " + x + " but found " + y);
                return true;
            }
            else if (x is TypeStack)
            {
                Assert(y is TypeStack, "expected type stack but found " + y.ToString());
                TypeStack tsx = x as TypeStack;
                TypeStack tsy = y as TypeStack;
                if (tsy.Count > tsx.Count) 
                    return false;
                for (int i = 0; i < tsy.Count; ++i)
                {
                    if (!IsSubType(tsx[i], tsy[i]))
                    {
                        return false;
                    }
                }
                return true;
            }
            else if (x is FxnType)
            {
                Assert(y is FxnType, "expected fxn type but found " + y.ToString());
                FxnType fx = x as FxnType;
                FxnType fy = y as FxnType;
                return (
                    IsSubType(fx.GetMainConsStack(), fy.GetAuxConsStack()) &&
                    IsSubType(fx.GetAuxConsStack(), fy.GetAuxConsStack()) &&
                    IsSubType(fx.GetMainProdStack(), fy.GetMainProdStack()) &&
                    IsSubType(fx.GetAuxProdStack(), fy.GetAuxProdStack()));
            }
            else
            {
                Assert(false, "unhandled subtype scenario for " + x.ToString() + " and " + y.ToString());
                return false;
            }
        }
        private CatType GetBetterConstraint(CatType x, CatType y)
        {
            if (x == y)
            {
                return x;
            }
            else if (x == null)
            {
                return y; 
            }
            else if (y == null)
            {
                return x;
            }
            else if (x.IsTypeVar())
            {
                if (y.IsTypeVar())
                {
                    // When there are two type variables, the preferred is the first one which occurs in the 
                    // constraint list.
                    TypeVarDecl ret = null;
                    foreach (TypeVarDecl t in mConstraints.Keys)
                    {
                        if ((ret == null) && (t == x.AsDecl() || t == y.AsDecl()))
                            ret = t;
                    }
                    Assert(ret != null, "could not find constraints " + x.ToString() + " or " + y.ToString());
                    return ret;
                }
                else
                {
                    return y;
                }
            }
            else if (y.IsTypeVar())
            {
                // We know x is not a TypeVar because it failed the previous type.
                return x;
            }
            else if (IsSubType(x, y))
            {
                return x;
            }
            else if (IsSubType(y, x))
            {
                return y;
            }
            else
            {
                throw new Exception("neither " + x.ToString() + " or " + y.ToString() + " are subtypes of each other");
            }
        }

        /// <summary>
        /// Goes through all of the constraints removing redundant constraints, choosing 
        /// the more specific or earlier constraint where multiple constraints conflict.
        /// It is important that this is called before any call to Resolve. 
        /// </summary>
        public void MergeAllConstraints()
        {
            Assert(!mbConstraintsMerged, "constraints are already merged");
            TypeVarsToValues();
            ResolveConstraintVars();
            ChooseBestVariables();
            mbConstraintsMerged = true;
        }

        private bool ReplaceAllInstances(TypeStack stk, TypeVarDecl decl, CatType t)
        {
            bool bRet = false;
            for (int i = stk.Count - 1; i >= 0; --i)
            {
                CatType u = stk[i];
                if (u.IsTypeVar())
                {
                    // NOTE: using TypeEq doesn't work because the type variable indexes
                    // weren't created by this point. This concerns me, and I am not sure
                    // that comparing pointers is sufficient. 
                    if (u.AsDecl() == decl)
                    {
                        stk.ReplaceAndExpand(i, t);
                        bRet = true;
                    }
                }
            }
            return bRet;
        }

        private bool ReplaceAllInstances(TypeVarDecl decl, CatType t)
        {
            bool bRet = false;
            foreach (TypeVarDecl key in mConstraints.Keys)
            {
                List<CatType> list = mConstraints[key];
                for (int i=list.Count - 1; i >= 0; --i)
                {
                    if (list[i].IsTypeVar())
                    {
                        TypeVarDecl tmp = list[i].AsDecl();
                        // NOTE: using TypeEq doesn't work because the type variable indexes
                        // weren't created by this point. This concerns me, and I am not sure
                        // that comparing pointers is sufficient. 
                        if (tmp == decl)
                        {
                            list[i] = t;
                            bRet = true;
                        }
                    }
                    else if (list[i] is TypeStack)
                    {
                        bRet = bRet || ReplaceAllInstances(list[i] as TypeStack, decl, t);
                    }
                }
            }
            return bRet;
        }

        /// <summary>
        /// Replace all instances of a TypeVar in the constraints 
        /// with concrete values. 
        /// </summary>
        private void TypeVarsToValues()
        {
            bool bReplacementMade = true;
            {
                while (bReplacementMade)
                {
                    bReplacementMade = false;
                    foreach (TypeVarDecl key in mConstraints.Keys)
                    {
                        CatType u = GetConcreteConstraint(key);
                        if (u != null)
                        {
                            if (ReplaceAllInstances(key, u))
                                bReplacementMade = true;
                        }
                    }
                }
            }

        }

        /// <summary>
        /// Given that each type var has a list of constraints, we need to choose the most 
        /// appropriate one from the list. 
        /// </summary>
        private void ChooseBestVariables()
        {
            foreach (CatType t in mConstraints.Keys)
            {
                List<CatType> list = mConstraints[t.AsDecl()];
                CatType best = null;
                foreach (CatType u in list)
                    best = GetBetterConstraint(best, u);
                list.Clear();
                list.Add(best);
            }
        }

        /// <summary>
        /// Resolves all of the constraints within a particular FxnType.
        /// </summary>
        public FxnType ResolveFxn(FxnType f)
        {
            Assert(mbConstraintsMerged, "missing call to MergeAllConstraints");
            FxnType ret = f;
            Resolve(ret);
            
            // Sometimes I think that a call to ret.Initialize should be here
            // but it causes a lot of problems in the type inference algorithm
            // ret.Initialize();

            return ret;
        }

        #endregion
    }
}