TwoThreeTree.n

/*
 * Copyright (c) 2009-2009  rampelstinskin@gmail.com
 * All rights reserved.
 *
 * Redistribution and use in source and binary forms, with or without
 * modification, are permitted provided that the following conditions
 * are met:
 *    1. Redistributions of source code must retain the above copyright
 *       notice, this list of conditions and the following disclaimer.
 *    2. Redistributions in binary form must reproduce the above copyright
 *       notice, this list of conditions and the following disclaimer in the
 *       documentation and/or other materials provided with the distribution.
 *    3. The name of rampelstinskin may not be used to endorse or promote
 *       products derived from this software without specific prior
 *       written permission.
 * 
 * THIS SOFTWARE IS PROVIDED BY THE AUTHOR ``AS IS'' AND ANY EXPRESS OR
 * IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED WARRANTIES
 * OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE DISCLAIMED. IN
 * NO EVENT SHALL THE AUTHOR BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL,
 * SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED
 * TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR
 * PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF
 * LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING
 * NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS
 * SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
 */

using System;
using SC = System.Collections;
using SCG = System.Collections.Generic;
using Nemerle.Text;
using Nemerle.Utility;

namespace Nemerle.Collections
{
  internal module TwoThreeTree
  {
    public enum InsertOptions
    {
    | Replace
    | ThrowIfDuplicate
    //| IgnoreDuplicate
    }

    public enum DeleteOptions
    {
    | IgnoreMissed
    | ThrowIfMissed
    }

    public variant Node[T]
    {
      | Leaf
      | Node2 { low : Node[T]; value1 : T; hi : Node[T] }
      | Node3 { low : Node[T]; value1 : T; mid : Node[T]; value2 : T; hi : Node[T] }

      public Insert(value : T, cmp : T * T -> int, options : InsertOptions, size : ref int) : Node[T]
      {
        mutable retL;
        mutable retV;
        mutable retH;
        mutable replaced = false;
        def eq(options, value)
        {
          match (options : InsertOptions)
          {
          | Replace => true
          //| IgnoreDuplicate => ();
          | ThrowIfDuplicate => throw ArgumentException($"\"$value\" already in the tree");
          | _                => throw ArgumentException($<#The parameter options has invalid value: "$options"#>, "options");
          }
        }
        def insert(_) : bool
        {
        | Leaf =>
          retL = Leaf();
          retV = value;
          retH = Leaf();
          false;
        | Node2 as node =>
          match (cmp(value, node.value1))
          {
          | c when c < 0 =>
            if (insert(node.low))
            {
              retL = Node2(retL, node.value1, node.hi);
              true;
            }
            else
            {
              retL = Node3(retL, retV, retH, node.value1, node.hi);
              true;
            }
          | c when c > 0 =>
            if (insert(node.hi))
            {
              retL = Node2(node.low, node.value1, retL);
              true;
            }
            else
            {
              retL = Node3(node.low, node.value1, retL, retV, retH);
              true;
            }
          | _ =>
            replaced = eq(options, value);
            retL = Node2(node.low, value, node.hi);
            true;
          }
        | Node3 as node =>
          match (cmp(value, node.value1))
          {
          | c when c < 0 =>
            if (insert(node.low))
            {
              retL = Node3(retL, node.value1, node.mid, node.value2, node.hi);
              true;
            }
            else
            {
              retL = Node2(retL, retV, retH);
              retV = node.value1;
              retH = Node2(node.mid, node.value2, node.hi);
              false;
            }
          | c when c > 0 =>
            match (cmp(value, node.value2))
            {
            | c when c < 0 =>
              if (insert(node.mid))
              {
                retL = Node3(node.low, node.value1, retL, node.value2, node.hi);
                true;
              }
              else
              {
                retL = Node2(node.low, node.value1, retL);
                //retV = retV;
                retH = Node2(retH, node.value2, node.hi);
                false;
              }
            | c when c > 0 =>
              if (insert(node.hi))
              {
                retL = Node3(node.low, node.value1, node.mid, node.value2, retL);
                true;
              }
              else
              {
                retH = Node2(retL, retV, retH);
                retL = Node2(node.low, node.value1, node.mid);
                retV = node.value2;
                false;
              }
            | _ =>
              replaced = eq(options, value);
              retL = Node3(node.low, node.value1, node.mid, value, node.hi);
              true;
            }
          | _ =>
            replaced = eq(options, value);
            retL = Node3(node.low, value, node.mid, node.value2, node.hi);
            true;
          }
        }
        def ret = if (insert(this))
          retL;
        else
          Node2(retL, retV, retH);
        unless (replaced)
          ++size;
        ret;
      }

      private abstract class CmpBase
      {
        public abstract Cmp21(node : Node[T].Node2) : int;
        public abstract Cmp31(node : Node[T].Node3) : int;
        public abstract Cmp32(node : Node[T].Node3) : int;
      }

      [Record]
      private class CmpValue : CmpBase
      {
        private _value : T;
        private _cmp : T * T -> int;

        public override Cmp21(node : Node[T].Node2) : int { _cmp(_value, node.value1) }
        public override Cmp31(node : Node[T].Node3) : int { _cmp(_value, node.value1) }
        public override Cmp32(node : Node[T].Node3) : int { _cmp(_value, node.value2) }
      }

      private class CmpReplacement : CmpBase
      {
        public override Cmp21(node : Node[T].Node2) : int { match (node.low) { | Leaf => 0 | _ => -1 } }
        public override Cmp31(node : Node[T].Node3) : int { match (node.low) { | Leaf => 0 | _ => -1 } }
        public override Cmp32(node : Node[T].Node3) : int { match (node.low) { | Leaf => 0 | _ => -1 } }
      }
      private static _cmpReplacement : CmpBase = CmpReplacement();

      public Delete(value : T, cmp : T * T -> int, options : DeleteOptions, size : ref int) : Node[T]
      {
        mutable retValue : T;
        mutable retNode : Node[T];
        mutable isMissed = false;
        def missed(node : Node[T])
        {
          match (options)
          {
          | ThrowIfMissed => throw ArgumentException($"\"$value\" not in the tree");
          | IgnoreMissed =>
            isMissed = true;
            retNode = node;
            false;
          }
        }
        //return isEmptyNode
        def delete(node : Node[T], cmp : CmpBase) : bool
        {
          match (node)
          {
          | Leaf as node => missed(node);
          | Node2 as node =>
            def balanceLow(isEmpty, value, node)
            {
              if (isEmpty)
              {
                match (node.hi)
                {
                | Node2 as child =>
                  retNode = Node3(retNode, value, child.low, child.value1, child.hi);
                  true;
                | Node3 as child =>
                  retNode = Node2(Node2(retNode, value, child.low), child.value1, Node2(child.mid, child.value2, child.hi));
                  false;
                | Leaf => assert(false);
                }
              }
              else
              {
                retNode = Node2(retNode, value, node.hi);
                false;
              }
            }
            def balanceHi(isEmpty, value, node)
            {
              if (isEmpty)
              {
                match (node.low)
                {
                | Node2 as child =>
                  retNode = Node3(child.low, child.value1, child.hi, value, retNode);
                  true;
                | Node3 as child =>
                  retNode = Node2(Node2(child.low, child.value1, child.mid), child.value2, Node2(child.hi, value, retNode));
                  false;
                | Leaf => assert(false);
                }
              }
              else
              {
                retNode = Node2(node.low, value, retNode);
                false;
              }
            }
            match (cmp.Cmp21(node))
            {
            | c when c < 0 =>
              balanceLow(delete(node.low, cmp), node.value1, node);
            | c when c > 0 =>
              balanceHi(delete(node.hi, cmp), node.value1, node);
            | _ =>
              match (node.low)
              {
              | Leaf =>
                retValue = node.value1;
                retNode = Leaf();
                true;
              | _ =>
                balanceHi(delete(node.hi, _cmpReplacement), retValue, node);
              }
            }
          | Node3 as node =>
            def balanceLow(isEmpty, value, node)
            {
              if (isEmpty)
              {
                match (node.mid)
                {
                | Node2 as child =>
                  retNode = Node2(Node3(retNode, value, child.low, child.value1, child.hi), node.value2, node.hi);
                  false;
                | Node3 as child =>
                  retNode = Node3(Node2(retNode, value, child.low), child.value1, Node2(child.mid, child.value2, child.hi), node.value2, node.hi);
                  false;
                | Leaf => assert(false);
                }
              }
              else
              {
                retNode = Node3(retNode, value, node.mid, node.value2, node.hi);
                false;
              }
            }
            def balanceMid(isEmpty, value, node)
            {
              if (isEmpty)
              {
                match (node.low)
                {
                | Node2 as child =>
                  retNode = Node2(Node3(child.low, child.value1, child.hi, value, retNode), node.value2, node.hi);
                  false;
                | Node3 as child =>
                  retNode = Node3(Node2(child.low, child.value1, child.mid), child.value2, Node2(child.hi, value, retNode), node.value2, node.hi);
                  false;
                | Leaf => assert(false);
                }
              }
              else
              {
                retNode = Node3(node.low, value, retNode, node.value2, node.hi);
                false;
              }
            }
            def balanceHi(isEmpty, value, node)
            {
              if (isEmpty)
              {
                match (node.mid)
                {
                | Node2 as child =>
                  retNode = Node2(node.low, node.value1, Node3(child.low, child.value1, child.hi, value, retNode));
                  false;
                | Node3 as child =>
                  retNode = Node3(node.low, node.value1, Node2(child.low, child.value1, child.mid), child.value2, Node2(child.hi, value, retNode));
                  false;
                | Leaf => assert(false);
                }
              }
              else
              {
                retNode = Node3(node.low, node.value1, node.mid, value, retNode);
                false;
              }
            }
            match (cmp.Cmp31(node))
            {
            | c when c < 0 =>
              balanceLow(delete(node.low, cmp), node.value1, node);
            | c when c > 0 =>
              match (cmp.Cmp32(node))
              {
              | c when c < 0 =>
                balanceMid(delete(node.mid, cmp), node.value1, node);
              | c when c > 0 =>
                balanceHi(delete(node.hi, cmp), node.value2, node);
              | _ =>
                match (node.low)
                {
                | Leaf =>
                  retValue = node.value2;
                  retNode = Node2(Leaf(), node.value1, Leaf());
                  false;
                | _ =>
                  balanceHi(delete(node.hi, _cmpReplacement), retValue, node);
                }
              }
            | _ =>
              match (node.low)
              {
              | Leaf =>
                retValue = node.value1;
                retNode = Node2(Leaf(), node.value2, Leaf());
                false;
              | _ =>
                balanceMid(delete(node.mid, _cmpReplacement), retValue, node);
              }
            }
          }
        }
        _ = delete(this, CmpValue(value, cmp));
        if (isMissed)
        {
          this;
        }
        else
        {
          --size;
          retNode;
        }
      }

      public Filter(cmp : T * T -> int, fn : T -> bool) : int * Node[T]
      {
        mutable size = 0;
        def filterFn(value, tree)
        {
          if (fn(value))
            tree.Insert(value, cmp, InsertOptions.ThrowIfDuplicate, ref size)
          else
            tree
        }
        (size, this.Fold(Node[T].Leaf(), filterFn));
      }

      public Partition(cmp : T * T -> int, fn : T -> bool) : int * Node[T] * int * Node[T]
      {
        mutable ysize = 0;
        mutable ytree = Node[T].Leaf();
        mutable nsize = 0;
        mutable ntree = Node[T].Leaf();
        def partitionFn(value, dummy)
        {
          if (fn(value))
            ytree = ytree.Insert(value, cmp, InsertOptions.ThrowIfDuplicate, ref ysize);
          else
            ntree = ntree.Insert(value, cmp, InsertOptions.ThrowIfDuplicate, ref nsize);
          dummy;
        }
        _ = this.Fold(null, partitionFn);
        (ysize, ytree, nsize, ntree);
      }

      public TryGet(value : T, cmp : T * T -> int, outValue : out T) : bool
      {
        def tryGet(node)
        {
        | Leaf => false
        | Node2 as node =>
          match (cmp(value, node.value1))
          {
          | c when c < 0 => tryGet(node.low);
          | c when c > 0 => tryGet(node.hi);
          | _ =>
            outValue = node.value1;
            true;
          }
        | Node3 as node =>
          match (cmp(value, node.value1))
          {
          | c when c < 0 => tryGet(node.low);
          | c when c > 0 =>
            match (cmp(value, node.value2))
            {
            | c when c < 0 => tryGet(node.mid);
            | c when c > 0 => tryGet(node.hi);
            | _ =>
              outValue = node.value2;
              true;
            }
          | _ =>
            outValue = node.value1;
            true;
          }
        }
        outValue = value;
        tryGet(this);
      }

      public Enumerate() : SCG.IEnumerable[T]
      {
        ToArray();
      }

      public Fold[TAccumulator](acc : TAccumulator, fn : T * TAccumulator -> TAccumulator) : TAccumulator
      {
        FoldLeft(acc, fn);
      }

      public FoldLeft[TAccumulator](acc : TAccumulator, fn : T * TAccumulator -> TAccumulator) : TAccumulator
      {
        match (this)
        {
        | Leaf => acc;
        | Node2 as node =>
          def acc = node.low.FoldLeft(acc, fn);
          def acc = fn(node.value1, acc);
          def acc = node.hi.FoldLeft(acc, fn);
          acc;
        | Node3 as node =>
          def acc = node.low.FoldLeft(acc, fn);
          def acc = fn(node.value1, acc);
          def acc = node.mid.FoldLeft(acc, fn);
          def acc = fn(node.value2, acc);
          def acc = node.hi.FoldLeft(acc, fn);
          acc;
        }
      }

      public FoldRight[TAccumulator](acc : TAccumulator, fn : T * TAccumulator -> TAccumulator) : TAccumulator
      {
        match (this)
        {
        | Leaf => acc;
        | Node2 as node =>
          def acc = node.hi.FoldRight(acc, fn);
          def acc = fn(node.value1, acc);
          def acc = node.low.FoldRight(acc, fn);
          acc;
        | Node3 as node =>
          def acc = node.hi.FoldRight(acc, fn);
          def acc = fn(node.value2, acc);
          def acc = node.mid.FoldRight(acc, fn);
          def acc = fn(node.value1, acc);
          def acc = node.low.FoldRight(acc, fn);
          acc;
        }
      }

      public ForAll[TAccumulator](fn : T -> bool) : bool
      {
        match (this)
        {
        | Leaf => true;
        | Node2 as node => fn(node.value1) && node.low.ForAll(fn) && node.hi.ForAll(fn);
        | Node3 as node => fn(node.value1) && fn(node.value2) && node.low.ForAll(fn) && node.mid.ForAll(fn) && node.hi.ForAll(fn);
        }
      }

      public Exists[TAccumulator](fn : T -> bool) : bool
      {
        match (this)
        {
        | Leaf => false;
        | Node2 as node => fn(node.value1) || node.low.Exists(fn) || node.hi.Exists(fn);
        | Node3 as node => fn(node.value1) || fn(node.value2) || node.low.Exists(fn) || node.mid.Exists(fn) || node.hi.Exists(fn);
        }
      }

      public Count : int
      {
        get
        {
          match (this)
          {
          | Leaf          => 0;
          | Node2 as node => 1 + node.low.Count + node.hi.Count;
          | Node3 as node => 2 + node.low.Count + node.mid.Count + node.hi.Count;
          }
        }
      }

      public Sum[TAccumulator](size : ref int, that : Node[T], cmp : T * T -> int) : Node[T]
      {
        mutable sz = size;
        def res = that.Fold(this, (value, tree) => tree.Insert(value, cmp, InsertOptions.Replace, ref sz));
        size = sz;
        res;
      }

      public Subtract[TAccumulator](size : ref int, that : Node[T], cmp : T * T -> int) : Node[T]
      {
        mutable sz = size;
        def res = that.Fold(this, (value, tree) => tree.Delete(value, cmp, DeleteOptions.IgnoreMissed, ref sz));
        size = sz;
        res;
      }

      public Intersect[TAccumulator](size : ref int, that : Node[T], cmp : T * T -> int) : Node[T]
      {
        mutable sz = 0;
        mutable outValue;
        def res = that.Fold
          (Node[T].Leaf(), (value, tree) =>
            if (this.TryGet(value, cmp, out outValue))
              tree.Insert(value, cmp, InsertOptions.ThrowIfDuplicate, ref sz)
            else
              tree
          );
        size = sz;
        res;
      }

      public Xor[TAccumulator](size : ref int, that : Node[T], cmp : T * T -> int) : Node[T]
      {
        mutable sz = 0;
        mutable outValue;
        def tree = Node[T].Leaf();
        def tree = that.Fold
          (tree, (value, tree) =>
            if (this.TryGet(value, cmp, out outValue))
              tree
            else
              tree.Insert(value, cmp, InsertOptions.ThrowIfDuplicate, ref sz)
          );
        def tree = this.Fold
          (tree, (value, tree) =>
            if (that.TryGet(value, cmp, out outValue))
              tree
            else
              tree.Insert(value, cmp, InsertOptions.ThrowIfDuplicate, ref sz)
          );
        size = sz;
        tree
      }

      public ToArray() : array[T]
      {
        FoldLeft(SCG.List(), (value, lst) => { lst.Add(value); lst}).ToArray();
      }

      public ToList() : list[T]
      {
        FoldRight([], _ :: _);
      }

      public MapToArray[U](fn : T -> U) : array[U]
      {
        FoldLeft(SCG.List(), (value, lst) => { lst.Add(fn(value)); lst}).ToArray();
      }

      public MapToList[U](fn : T -> U) : list[U]
      {
        FoldRight([], (value, acc) => fn(value) :: acc);
      }

      public ToString(first : string, delim : string, last : string) : string
      {
        def sb = System.Text.StringBuilder().Append(first);
        def sb = FoldLeft(sb, (value, sb) => sb.Append(value : object).Append(delim) );
        when (sb.Length > first.Length)
          sb.Length -= delim.Length;
        _ = sb.Append(last);
        sb.ToString();
      }

      public override GetHashCode() : int
      {
        FoldLeft(0, (val, hash) => unchecked { (hash ^ val.GetHashCode()) * 5 + hash } )
      }
    }
  }
}