container/mutable_tree_map.fz

# This file is part of the Fuzion language implementation.
#
# The Fuzion language implementation is free software: you can redistribute it
# and/or modify it under the terms of the GNU General Public License as published
# by the Free Software Foundation, version 3 of the License.
#
# The Fuzion language implementation is distributed in the hope that it will be
# useful, but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the GNU General Public
# License for more details.
#
# You should have received a copy of the GNU General Public License along with The
# Fuzion language implementation.  If not, see <https://www.gnu.org/licenses/>.


# -----------------------------------------------------------------------
#
#  Tokiwa Software GmbH, Germany
#
#  Source code of Fuzion standard library feature mutable_tree_map
#
# -----------------------------------------------------------------------


# mutable_tree_map -- a mutable map using an AVL tree
#
private:public mutable_tree_map(public LM  type : mutate,
                                public KEY type : property.orderable,
                                public VAL type                      ) : Mutable_Map KEY VAL
is


  # the root entry of this map
  #
  # mutable because this might change, for example when the key stored at the root
  # is removed.
  #
  # if this map is entry, this is nil
  #
  root := LM.env.new (option Entry) nil


  # returns the size of the map, i.e. the number of elements it contains
  #
  public redef size i32 =>
    fold i32 0 ((i, _) -> i + 1)


  # get a sequence of all key/value pairs in this map
  #
  public redef items Sequence (tuple KEY VAL) =>
    match root.get
      e Entry => e.items
      nil => (list (tuple KEY VAL)).empty


  # freeze the map, such that it is no longer mutable afterwards
  #
  public freeze unit =>
    if root.open
      root.close

    fold unit ((_, x) ->
      if x.left.open
        x.left.close

      if x.right.open
        x.right.close

      unit)


  # get the value stored in this map at key k, nil if k is not
  # contained in this map
  #
  public redef get(k KEY) option VAL =>
    root.get.bind (.get k)


  # add the mapping k -> v as a new entry to this map
  #
  public redef put(k KEY, v VAL) unit =>
    _ := replace k v

  # add the mapping k -> v as a new entry to this map
  #
  # returns the value that k previously mapped to, or nil if
  # k was not yet contained in this map
  #
  public replace(k KEY, v VAL) option VAL =>
    # helper feature to add a mapping to this map. this feature
    # additionally takes the node we are currently working at, and
    # also returns any new node, or the reference to the existing
    # node that was worked on.
    #
    # in this step, the helper feature to actually add the mapping
    # is called first. then the AVL rebalancing is done
    #
    put_recursively(node option Entry) tuple (option Entry) (option VAL) =>
      new_node, old_val := insert_or_modify_entries node
      (rebalance new_node, old_val)


    # helper feature to add a mapping to this map. this feature
    # additionally takes the node we are currently working at, and
    # also returns any new node, or the reference to the existing
    # node that was worked on.
    #
    # in this step, the actual addition of the mapping to the binary
    # tree is done, but this might violate the AVL invariants.
    #
    insert_or_modify_entries(node option Entry) tuple (option Entry) (option VAL) =>
      match node
        nil =>
          new_node := Entry k v
          (option new_node, option VAL nil)
        e Entry =>
          if k < e.key
            node0, old_val := put_recursively e.left.get
            e.left <- node0
            (option e, old_val)
          else if e.key < k
            node0, old_val := put_recursively e.right.get
            e.right <- node0
            (option e, old_val)
          else
            old_val := e.val
            new_node := Entry k v
            new_node.left <- e.left.get
            new_node.right <- e.right.get
            (option new_node, option old_val)


    new_root, old := put_recursively root.get
    root <- new_root
    old


  # remove the mapping from k to some value from this map
  #
  # returns the value that k previously mapped to, or nil if
  # no mapping was actually removed
  #
  public redef remove(key KEY) option VAL =>
    # helper feature to remove a mapping from this map. this feature
    # additionally takes the node we are currently working at, and
    # also returns the reference to the node that was worked on.
    #
    # in this step, the helper feature to actually remove the
    # mapping is called first. then the AVL rebalancing is done
    #
    remove_recursively(k KEY, node option Entry) tuple (option Entry) (option VAL) =>
      new_node, old_val := remove_or_modify_entries k node
      (rebalance new_node, old_val)


    # search the subtree whose root is the given node for its minimal
    # node and return it
    #
    # minimal here means the node with the smallest key, by the given
    # ordering of the keys
    #
    minimum(node option Entry) option Entry =>
      node.bind e->
        e.left.get ? nil => e
                   | l Entry => minimum l


    # helper feature to remove a mapping from this map. this feature
    # additionally takes the node we are currently working at, and
    # also returns the reference to the node that was worked on.
    #
    # in this step, the actual removal of the mapping from the binary
    # tree is done, but this might violate the AVL invariants.
    #
    remove_or_modify_entries(k KEY, node option Entry) tuple (option Entry) (option VAL) =>
      match node
        nil => (option Entry nil, option VAL nil)
        e Entry =>
          if k < e.key
            n, old_val := remove_recursively k e.left.get
            e.left <- n
            (option e, old_val)
          else if e.key < k
            n, old_val := remove_recursively k e.right.get
            e.right <- n
            (option e, old_val)
          else
            old_val := e.val
            match e.left.get
              nil => (e.right.get, option old_val)
              l Entry =>
                match e.right.get
                  nil   => (option l, option old_val)
                  Entry =>
                    m := minimum e.right.get

                    new_node := Entry m.get.key m.get.val
                    new_node.left <- l

                    nr, old_val0 := remove_recursively m.get.key e.right.get
                    new_node.right <- nr
                    (option new_node, option old_val0)


    new_root, old := remove_recursively key root.get
    root <- new_root
    old


  # create an immutable map from this
  #
  public redef as_map container.Map KEY VAL =>
    (container.ps_map KEY VAL).new keys.as_array values.as_array


  # rebalances a tree whose AVL invariants might be violated
  #
  # this determines the balance factor of the given node and applies
  # the appropriate rotations
  #
  rebalance(node option Entry) =>
    # returns the height of the subtree whose root is the given
    # node, or -1 if an empty subtree is given
    #
    height(n option Entry) =>
      match n
        nil => -1
        e Entry => e.height.get


    # returns the (AVL) balance factor of the given node, or
    # 0 if a nil node is given
    #
    balance_factor(n option Entry) =>
      match n
        nil => 0
        e Entry => (height e.right.get) - (height e.left.get)


    # recalculates and updates the heights of the subtrees in the
    # subtree whose root is the given node
    #
    fix_height(n option Entry) =>
      _ := n.bind e->
        lh := height e.left.get
        rh := height e.right.get

        e.height <- (if lh > rh then lh else rh) + 1


    # rotate right at the given node
    #
    rotate_right(n option Entry) =>
      # because this feature is only called when the tree is out of balance,
      # i.e. the left subtree has more nodes than the right one, we can safely
      # assume here that node and node.left are not nil.
      l := n.get.left.get

      n.get.left <- l.get.right.get
      l.get.right <- n

      fix_height n
      fix_height l

      l


    # rotate left at the given node
    #
    rotate_left(n option Entry) =>
      # because this feature is only called when the tree is out of balance,
      # i.e. the right subtree has more nodes than the left one, we can safely
      # assume here that node and node.right are not nil.
      r := n.get.right.get

      n.get.right <- r.get.left.get
      r.get.left <- n

      fix_height n
      fix_height r

      r


    # rebalance is called when nodes might have been inserted or deleted.
    # this means that the heights of the nodes potentially changed, thus
    # we need to recalculate them.
    fix_height node

    bf := balance_factor node

    if bf < -1
      # we can safely do node.get here because if node was empty,
      # its balance factor would be 0.
      if balance_factor node.get.left.get <= 0
        rotate_right node
      else
        node.get.left <- rotate_left node.get.left.get
        rotate_right node
    else if bf > 1
      # we can safely do node.get here because if node was empty,
      # its balance factor would be 0.
      if balance_factor node.get.right.get >= 0
        rotate_left node
      else
        node.get.right <- rotate_right node.get.right.get
        rotate_left node
    else
      # do nothing
      node


  # perform an in-order traversal of the tree and process the entries
  # encountered using the initial value init and the combinator feature
  # f. the latter takes the last result of the computation and the node
  # currently visited and combines this information in some way.
  #
  fold(B type, init B, f (B, Entry) -> B) =>
    fold0(init B, node option Entry) =>
      node ? nil => fold0.this.init
           | n Entry => fold0 (f (fold0 fold0.this.init n.left.get) n) n.right.get

    fold0 init root.get


  # returns an empty tree of elements of type A.
  #
  public fixed redef type.empty container.mutable_tree_map LM KEY VAL =>
    container.mutable_tree_map LM KEY VAL


  # returns a tree of elements of type A that contains just the element a.
  #
  public type.from_key_value(k KEY, v VAL) container.mutable_tree_map LM KEY VAL =>
    new_map := empty
    new_map.put k v

    new_map


  # initialize a map from a Sequence of key value tuples
  #
  public type.from_sequence(kvs Sequence (tuple KEY VAL)) container.mutable_tree_map LM KEY VAL =>
    from_sequence kvs false


  # initialize a map from a Sequence of key value tuples
  #
  # if the freeze argument is true, then the map is frozen
  # after being populated with the entries from the Sequence.
  #
  public type.from_sequence(kvs Sequence (tuple KEY VAL), freeze bool) container.mutable_tree_map LM KEY VAL =>
    new_map := empty
    kvs.for_each x->
      new_map.put x.0 x.1

    if freeze
      new_map.freeze

    new_map