1/*
2 * Copyright (C) 2011 The Guava Authors
3 *
4 * Licensed under the Apache License, Version 2.0 (the "License"); you may not use this file except
5 * in compliance with the License. You may obtain a copy of the License at
6 *
7 * http://www.apache.org/licenses/LICENSE-2.0
8 *
9 * Unless required by applicable law or agreed to in writing, software distributed under the
10 * License is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either
11 * express or implied. See the License for the specific language governing permissions and
12 * limitations under the License.
13 */
14
15package com.google.common.collect;
16
17import static com.google.common.base.Preconditions.checkNotNull;
18import static com.google.common.collect.BstSide.LEFT;
19import static com.google.common.collect.BstSide.RIGHT;
20
21import com.google.common.annotations.GwtCompatible;
22
23import java.util.Comparator;
24
25import javax.annotation.Nullable;
26
27/**
28 * Tools to perform single-key queries and mutations in binary search trees.
29 *
30 * @author Louis Wasserman
31 */
32@GwtCompatible
33final class BstOperations {
34  private BstOperations() {}
35
36  /**
37   * Returns the node with key {@code key} in {@code tree}, if any.
38   */
39  @Nullable
40  public static <K, N extends BstNode<K, N>> N seek(
41      Comparator<? super K> comparator, @Nullable N tree, @Nullable K key) {
42    checkNotNull(comparator);
43    if (tree == null) {
44      return null;
45    }
46    int cmp = comparator.compare(key, tree.getKey());
47    if (cmp == 0) {
48      return tree;
49    } else {
50      BstSide side = (cmp < 0) ? LEFT : RIGHT;
51      return seek(comparator, tree.childOrNull(side), key);
52    }
53  }
54
55  /**
56   * Returns the result of performing the mutation specified by {@code mutationRule} in {@code
57   * tree} at the location with key {@code key}.
58   *
59   * <ul>
60   * <li>If the returned {@link BstModificationResult} has type {@code IDENTITY}, the exact
61   * original tree is returned.
62   * <li>If the returned {@code BstModificationResult} has type {@code REBUILDING_CHANGE},
63   * the tree will be rebuilt with the node factory of the mutation rule, but not rebalanced.
64   * <li>If the returned {@code BstModificationResult} has type {@code REBALANCING_CHANGE},
65   * the tree will be rebalanced using the balance policy of the mutation rule.
66   * </ul>
67   */
68  public static <K, N extends BstNode<K, N>> BstMutationResult<K, N> mutate(
69      Comparator<? super K> comparator, BstMutationRule<K, N> mutationRule, @Nullable N tree,
70      @Nullable K key) {
71    checkNotNull(comparator);
72    checkNotNull(mutationRule);
73
74    if (tree != null) {
75      int cmp = comparator.compare(key, tree.getKey());
76      if (cmp != 0) {
77        BstSide side = (cmp < 0) ? LEFT : RIGHT;
78        BstMutationResult<K, N> mutation =
79            mutate(comparator, mutationRule, tree.childOrNull(side), key);
80        return mutation.lift(
81            tree, side, mutationRule.getNodeFactory(), mutationRule.getBalancePolicy());
82      }
83    }
84    return modify(tree, key, mutationRule);
85  }
86
87  /**
88   * Perform the local mutation at the tip of the specified path.
89   */
90  public static <K, N extends BstNode<K, N>> BstMutationResult<K, N> mutate(
91      BstInOrderPath<N> path, BstMutationRule<K, N> mutationRule) {
92    checkNotNull(path);
93    checkNotNull(mutationRule);
94    BstBalancePolicy<N> balancePolicy = mutationRule.getBalancePolicy();
95    BstNodeFactory<N> nodeFactory = mutationRule.getNodeFactory();
96    BstModifier<K, N> modifier = mutationRule.getModifier();
97
98    N target = path.getTip();
99    K key = target.getKey();
100    BstMutationResult<K, N> result = modify(target, key, mutationRule);
101    while (path.hasPrefix()) {
102      BstInOrderPath<N> prefix = path.getPrefix();
103      result = result.lift(prefix.getTip(), path.getSideOfExtension(), nodeFactory, balancePolicy);
104      path = prefix;
105    }
106    return result;
107  }
108
109  /**
110   * Perform the local mutation right here, at the specified node.
111   */
112  private static <K, N extends BstNode<K, N>> BstMutationResult<K, N> modify(
113      @Nullable N tree, K key, BstMutationRule<K, N> mutationRule) {
114    BstBalancePolicy<N> rebalancePolicy = mutationRule.getBalancePolicy();
115    BstNodeFactory<N> nodeFactory = mutationRule.getNodeFactory();
116    BstModifier<K, N> modifier = mutationRule.getModifier();
117
118    N originalRoot = tree;
119    N changedRoot;
120    N originalTarget = (tree == null) ? null : nodeFactory.createLeaf(tree);
121    BstModificationResult<N> modResult = modifier.modify(key, originalTarget);
122    N originalLeft = null;
123    N originalRight = null;
124    if (tree != null) {
125      originalLeft = tree.childOrNull(LEFT);
126      originalRight = tree.childOrNull(RIGHT);
127    }
128    switch (modResult.getType()) {
129      case IDENTITY:
130        changedRoot = tree;
131        break;
132      case REBUILDING_CHANGE:
133        if (modResult.getChangedTarget() != null) {
134          changedRoot =
135              nodeFactory.createNode(modResult.getChangedTarget(), originalLeft, originalRight);
136        } else if (tree == null) {
137          changedRoot = null;
138        } else {
139          throw new AssertionError(
140              "Modification result is a REBUILDING_CHANGE, but rebalancing required");
141        }
142        break;
143      case REBALANCING_CHANGE:
144        if (modResult.getChangedTarget() != null) {
145          changedRoot = rebalancePolicy.balance(
146              nodeFactory, modResult.getChangedTarget(), originalLeft, originalRight);
147        } else if (tree != null) {
148          changedRoot = rebalancePolicy.combine(nodeFactory, originalLeft, originalRight);
149        } else {
150          changedRoot = null;
151        }
152        break;
153      default:
154        throw new AssertionError();
155    }
156    return BstMutationResult.mutationResult(key, originalRoot, changedRoot, modResult);
157  }
158
159  /**
160   * Returns the result of removing the minimum element from the specified subtree.
161   */
162  public static <K, N extends BstNode<K, N>> BstMutationResult<K, N> extractMin(
163      N root, BstNodeFactory<N> nodeFactory, BstBalancePolicy<N> balancePolicy) {
164    checkNotNull(root);
165    checkNotNull(nodeFactory);
166    checkNotNull(balancePolicy);
167    if (root.hasChild(LEFT)) {
168      BstMutationResult<K, N> subResult =
169          extractMin(root.getChild(LEFT), nodeFactory, balancePolicy);
170      return subResult.lift(root, LEFT, nodeFactory, balancePolicy);
171    }
172    return BstMutationResult.mutationResult(
173        root.getKey(), root, root.childOrNull(RIGHT),
174        BstModificationResult.rebalancingChange(root, null));
175  }
176
177  /**
178   * Returns the result of removing the maximum element from the specified subtree.
179   */
180  public static <K, N extends BstNode<K, N>> BstMutationResult<K, N> extractMax(
181      N root, BstNodeFactory<N> nodeFactory, BstBalancePolicy<N> balancePolicy) {
182    checkNotNull(root);
183    checkNotNull(nodeFactory);
184    checkNotNull(balancePolicy);
185    if (root.hasChild(RIGHT)) {
186      BstMutationResult<K, N> subResult =
187          extractMax(root.getChild(RIGHT), nodeFactory, balancePolicy);
188      return subResult.lift(root, RIGHT, nodeFactory, balancePolicy);
189    }
190    return BstMutationResult.mutationResult(root.getKey(), root, root.childOrNull(LEFT),
191        BstModificationResult.rebalancingChange(root, null));
192  }
193
194  /**
195   * Inserts the specified entry into the tree as the minimum entry. Assumes that {@code
196   * entry.getKey()} is less than the key of all nodes in the subtree {@code root}.
197   */
198  public static <N extends BstNode<?, N>> N insertMin(@Nullable N root, N entry,
199      BstNodeFactory<N> nodeFactory, BstBalancePolicy<N> balancePolicy) {
200    checkNotNull(entry);
201    checkNotNull(nodeFactory);
202    checkNotNull(balancePolicy);
203    if (root == null) {
204      return nodeFactory.createLeaf(entry);
205    } else {
206      return balancePolicy.balance(nodeFactory, root,
207          insertMin(root.childOrNull(LEFT), entry, nodeFactory, balancePolicy),
208          root.childOrNull(RIGHT));
209    }
210  }
211
212  /**
213   * Inserts the specified entry into the tree as the maximum entry. Assumes that {@code
214   * entry.getKey()} is greater than the key of all nodes in the subtree {@code root}.
215   */
216  public static <N extends BstNode<?, N>> N insertMax(@Nullable N root, N entry,
217      BstNodeFactory<N> nodeFactory, BstBalancePolicy<N> balancePolicy) {
218    checkNotNull(entry);
219    checkNotNull(nodeFactory);
220    checkNotNull(balancePolicy);
221    if (root == null) {
222      return nodeFactory.createLeaf(entry);
223    } else {
224      return balancePolicy.balance(nodeFactory, root, root.childOrNull(LEFT),
225          insertMax(root.childOrNull(RIGHT), entry, nodeFactory, balancePolicy));
226    }
227  }
228}
229