firebase-ios-sdk/Firebase/Database/Core/Utilities/FImmutableTree.m

/*
 * Copyright 2017 Google
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 *      http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */

#import "FImmutableTree.h"
#import "FImmutableSortedDictionary.h"
#import "FPath.h"
#import "FUtilities.h"

@interface FImmutableTree ()
@property(nonatomic, strong, readwrite) id value;
/**
 * Maps NSString -> FImmutableTree<T>, where <T> is type of value.
 */
@property(nonatomic, strong, readwrite) FImmutableSortedDictionary *children;
@end

@implementation FImmutableTree
@synthesize value;
@synthesize children;

- (id)initWithValue:(id)aValue {
    self = [super init];
    if (self) {
        self.value = aValue;
        self.children = [FImmutableTree emptyChildren];
    }
    return self;
}

- (id)initWithValue:(id)aValue
           children:(FImmutableSortedDictionary *)childrenMap {
    self = [super init];
    if (self) {
        self.value = aValue;
        self.children = childrenMap;
    }
    return self;
}

+ (FImmutableSortedDictionary *)emptyChildren {
    static dispatch_once_t emptyChildrenToken;
    static FImmutableSortedDictionary *emptyChildren;
    dispatch_once(&emptyChildrenToken, ^{
      emptyChildren = [FImmutableSortedDictionary
          dictionaryWithComparator:[FUtilities stringComparator]];
    });
    return emptyChildren;
}

+ (FImmutableTree *)empty {
    static dispatch_once_t emptyImmutableTreeToken;
    static FImmutableTree *emptyTree = nil;
    dispatch_once(&emptyImmutableTreeToken, ^{
      emptyTree = [[FImmutableTree alloc] initWithValue:nil];
    });
    return emptyTree;
}

- (BOOL)isEmpty {
    return self.value == nil && [self.children isEmpty];
}

/**
 * Given a path and a predicate, return the first node and the path to that node
 * where the predicate returns true
 * // TODO Do a perf test. If we're creating a bunch of FTuplePathValue objects
 * on the way back out, it may be better to pass down a pathSoFar FPath
 */
- (FTuplePathValue *)findRootMostMatchingPath:(FPath *)relativePath
                                    predicate:(BOOL (^)(id value))predicate {
    if (self.value != nil && predicate(self.value)) {
        return [[FTuplePathValue alloc] initWithPath:[FPath empty]
                                               value:self.value];
    } else {
        if ([relativePath isEmpty]) {
            return nil;
        } else {
            NSString *front = [relativePath getFront];
            FImmutableTree *child = [self.children get:front];
            if (child != nil) {
                FTuplePathValue *childExistingPathAndValue =
                    [child findRootMostMatchingPath:[relativePath popFront]
                                          predicate:predicate];
                if (childExistingPathAndValue != nil) {
                    FPath *fullPath = [[[FPath alloc] initWith:front]
                        child:childExistingPathAndValue.path];
                    return [[FTuplePathValue alloc]
                        initWithPath:fullPath
                               value:childExistingPathAndValue.value];
                } else {
                    return nil;
                }
            } else {
                // No child matching path
                return nil;
            }
        }
    }
}

/**
 * Find, if it exists, the shortest subpath of the given path that points a
 * defined value in the tree
 */
- (FTuplePathValue *)findRootMostValueAndPath:(FPath *)relativePath {
    return [self findRootMostMatchingPath:relativePath
                                predicate:^BOOL(__unsafe_unretained id value) {
                                  return YES;
                                }];
}

- (id)rootMostValueOnPath:(FPath *)path {
    return [self rootMostValueOnPath:path
                            matching:^BOOL(id value) {
                              return YES;
                            }];
}

- (id)rootMostValueOnPath:(FPath *)path matching:(BOOL (^)(id))predicate {
    if (self.value != nil && predicate(self.value)) {
        return self.value;
    } else if (path.isEmpty) {
        return nil;
    } else {
        return [[self.children get:path.getFront]
            rootMostValueOnPath:[path popFront]
                       matching:predicate];
    }
}

- (id)leafMostValueOnPath:(FPath *)path {
    return [self leafMostValueOnPath:path
                            matching:^BOOL(id value) {
                              return YES;
                            }];
}

- (id)leafMostValueOnPath:(FPath *)relativePath
                 matching:(BOOL (^)(id))predicate {
    __block id currentValue = self.value;
    __block FImmutableTree *currentTree = self;
    [relativePath enumerateComponentsUsingBlock:^(NSString *key, BOOL *stop) {
      currentTree = [currentTree.children get:key];
      if (currentTree == nil) {
          *stop = YES;
      } else {
          id treeValue = currentTree.value;
          if (treeValue != nil && predicate(treeValue)) {
              currentValue = treeValue;
          }
      }
    }];
    return currentValue;
}

- (BOOL)containsValueMatching:(BOOL (^)(id))predicate {
    if (self.value != nil && predicate(self.value)) {
        return YES;
    } else {
        __block BOOL found = NO;
        [self.children enumerateKeysAndObjectsUsingBlock:^(
                           NSString *key, FImmutableTree *subtree, BOOL *stop) {
          found = [subtree containsValueMatching:predicate];
          if (found)
              *stop = YES;
        }];
        return found;
    }
}

- (FImmutableTree *)subtreeAtPath:(FPath *)relativePath {
    if ([relativePath isEmpty]) {
        return self;
    } else {
        NSString *front = [relativePath getFront];
        FImmutableTree *childTree = [self.children get:front];
        if (childTree != nil) {
            return [childTree subtreeAtPath:[relativePath popFront]];
        } else {
            return [FImmutableTree empty];
        }
    }
}

/**
 * Sets a value at the specified path
 */
- (FImmutableTree *)setValue:(id)newValue atPath:(FPath *)relativePath {
    if ([relativePath isEmpty]) {
        return [[FImmutableTree alloc] initWithValue:newValue
                                            children:self.children];
    } else {
        NSString *front = [relativePath getFront];
        FImmutableTree *child = [self.children get:front];
        if (child == nil) {
            child = [FImmutableTree empty];
        }
        FImmutableTree *newChild = [child setValue:newValue
                                            atPath:[relativePath popFront]];
        FImmutableSortedDictionary *newChildren =
            [self.children insertKey:front withValue:newChild];
        return [[FImmutableTree alloc] initWithValue:self.value
                                            children:newChildren];
    }
}

/**
 * Remove the value at the specified path
 */
- (FImmutableTree *)removeValueAtPath:(FPath *)relativePath {
    if ([relativePath isEmpty]) {
        if ([self.children isEmpty]) {
            return [FImmutableTree empty];
        } else {
            return [[FImmutableTree alloc] initWithValue:nil
                                                children:self.children];
        }
    } else {
        NSString *front = [relativePath getFront];
        FImmutableTree *child = [self.children get:front];
        if (child) {
            FImmutableTree *newChild =
                [child removeValueAtPath:[relativePath popFront]];
            FImmutableSortedDictionary *newChildren;
            if ([newChild isEmpty]) {
                newChildren = [self.children removeKey:front];
            } else {
                newChildren = [self.children insertKey:front
                                             withValue:newChild];
            }
            if (self.value == nil && [newChildren isEmpty]) {
                return [FImmutableTree empty];
            } else {
                return [[FImmutableTree alloc] initWithValue:self.value
                                                    children:newChildren];
            }
        } else {
            return self;
        }
    }
}

/**
 * Gets a value from the tree
 */
- (id)valueAtPath:(FPath *)relativePath {
    if ([relativePath isEmpty]) {
        return self.value;
    } else {
        NSString *front = [relativePath getFront];
        FImmutableTree *child = [self.children get:front];
        if (child) {
            return [child valueAtPath:[relativePath popFront]];
        } else {
            return nil;
        }
    }
}

/**
 * Replaces the subtree at the specified path with the given new tree
 */
- (FImmutableTree *)setTree:(FImmutableTree *)newTree
                     atPath:(FPath *)relativePath {
    if ([relativePath isEmpty]) {
        return newTree;
    } else {
        NSString *front = [relativePath getFront];
        FImmutableTree *child = [self.children get:front];
        if (child == nil) {
            child = [FImmutableTree empty];
        }
        FImmutableTree *newChild = [child setTree:newTree
                                           atPath:[relativePath popFront]];
        FImmutableSortedDictionary *newChildren;
        if ([newChild isEmpty]) {
            newChildren = [self.children removeKey:front];
        } else {
            newChildren = [self.children insertKey:front withValue:newChild];
        }
        return [[FImmutableTree alloc] initWithValue:self.value
                                            children:newChildren];
    }
}

/**
 * Performs a depth first fold on this tree. Transforms a tree into a single
 * value, given a function that operates on the path to a node, an optional
 * current value, and a map of the child names to folded subtrees
 */
- (id)foldWithBlock:(id (^)(FPath *path, id value,
                            NSDictionary *foldedChildren))block {
    return [self foldWithPathSoFar:[FPath empty] withBlock:block];
}

/**
 * Recursive helper for public facing foldWithBlock: method
 */
- (id)foldWithPathSoFar:(FPath *)pathSoFar
              withBlock:(id (^)(FPath *path, id value,
                                NSDictionary *foldedChildren))block {
    __block NSMutableDictionary *accum = [[NSMutableDictionary alloc] init];
    [self.children
        enumerateKeysAndObjectsUsingBlock:^(
            NSString *childKey, FImmutableTree *childTree, BOOL *stop) {
          accum[childKey] =
              [childTree foldWithPathSoFar:[pathSoFar childFromString:childKey]
                                 withBlock:block];
        }];
    return block(pathSoFar, self.value, accum);
}

/**
 * Find the first matching value on the given path. Return the result of
 * applying block to it.
 */
- (id)findOnPath:(FPath *)path
    andApplyBlock:(id (^)(FPath *path, id value))block {
    return [self findOnPath:path pathSoFar:[FPath empty] andApplyBlock:block];
}

- (id)findOnPath:(FPath *)pathToFollow
        pathSoFar:(FPath *)pathSoFar
    andApplyBlock:(id (^)(FPath *path, id value))block {
    id result = self.value ? block(pathSoFar, self.value) : nil;
    if (result != nil) {
        return result;
    } else {
        if ([pathToFollow isEmpty]) {
            return nil;
        } else {
            NSString *front = [pathToFollow getFront];
            FImmutableTree *nextChild = [self.children get:front];
            if (nextChild != nil) {
                return [nextChild findOnPath:[pathToFollow popFront]
                                   pathSoFar:[pathSoFar childFromString:front]
                               andApplyBlock:block];
            } else {
                return nil;
            }
        }
    }
}
/**
 * Call the block on each value along the path for as long as that function
 * returns true
 * @return The path to the deepest location inspected
 */
- (FPath *)forEachOnPath:(FPath *)path whileBlock:(BOOL (^)(FPath *, id))block {
    return [self forEachOnPath:path pathSoFar:[FPath empty] whileBlock:block];
}

- (FPath *)forEachOnPath:(FPath *)pathToFollow
               pathSoFar:(FPath *)pathSoFar
              whileBlock:(BOOL (^)(FPath *, id))block {
    if ([pathToFollow isEmpty]) {
        if (self.value) {
            block(pathSoFar, self.value);
        }
        return pathSoFar;
    } else {
        BOOL shouldContinue = YES;
        if (self.value) {
            shouldContinue = block(pathSoFar, self.value);
        }
        if (shouldContinue) {
            NSString *front = [pathToFollow getFront];
            FImmutableTree *nextChild = [self.children get:front];
            if (nextChild) {
                return
                    [nextChild forEachOnPath:[pathToFollow popFront]
                                   pathSoFar:[pathSoFar childFromString:front]
                                  whileBlock:block];
            } else {
                return pathSoFar;
            }
        } else {
            return pathSoFar;
        }
    }
}

- (FImmutableTree *)forEachOnPath:(FPath *)path
                     performBlock:(void (^)(FPath *path, id value))block {
    return [self forEachOnPath:path pathSoFar:[FPath empty] performBlock:block];
}

- (FImmutableTree *)forEachOnPath:(FPath *)pathToFollow
                        pathSoFar:(FPath *)pathSoFar
                     performBlock:(void (^)(FPath *path, id value))block {
    if ([pathToFollow isEmpty]) {
        return self;
    } else {
        if (self.value) {
            block(pathSoFar, self.value);
        }
        NSString *front = [pathToFollow getFront];
        FImmutableTree *nextChild = [self.children get:front];
        if (nextChild) {
            return [nextChild forEachOnPath:[pathToFollow popFront]
                                  pathSoFar:[pathSoFar childFromString:front]
                               performBlock:block];
        } else {
            return [FImmutableTree empty];
        }
    }
}
/**
 * Calls the given block for each node in the tree that has a value. Called in
 * depth-first order
 */
- (void)forEach:(void (^)(FPath *path, id value))block {
    [self forEachPathSoFar:[FPath empty] withBlock:block];
}

- (void)forEachPathSoFar:(FPath *)pathSoFar
               withBlock:(void (^)(FPath *path, id value))block {
    [self.children
        enumerateKeysAndObjectsUsingBlock:^(
            NSString *childKey, FImmutableTree *childTree, BOOL *stop) {
          [childTree forEachPathSoFar:[pathSoFar childFromString:childKey]
                            withBlock:block];
        }];
    if (self.value) {
        block(pathSoFar, self.value);
    }
}

- (void)forEachChild:(void (^)(NSString *childKey, id childValue))block {
    [self.children
        enumerateKeysAndObjectsUsingBlock:^(
            NSString *childKey, FImmutableTree *childTree, BOOL *stop) {
          if (childTree.value) {
              block(childKey, childTree.value);
          }
        }];
}

- (BOOL)isEqual:(id)object {
    if (![object isKindOfClass:[FImmutableTree class]]) {
        return NO;
    }
    FImmutableTree *other = (FImmutableTree *)object;
    return (self.value == other.value || [self.value isEqual:other.value]) &&
           [self.children isEqual:other.children];
}

- (NSUInteger)hash {
    return self.children.hash * 31 + [self.value hash];
}

- (NSString *)description {
    NSMutableString *string = [[NSMutableString alloc] init];
    [string appendString:@"FImmutableTree { value="];
    [string appendString:(self.value ? [self.value description] : @"<nil>")];
    [string appendString:@", children={"];
    [self.children
        enumerateKeysAndObjectsUsingBlock:^(
            NSString *childKey, FImmutableTree *childTree, BOOL *stop) {
          [string appendString:@" "];
          [string appendString:childKey];
          [string appendString:@"="];
          [string appendString:[childTree.value description]];
        }];
    [string appendString:@" } }"];
    return [NSString stringWithString:string];
}

- (NSString *)debugDescription {
    return [self description];
}

@end