From 901aaeac47c2275d0e5facbf9c5525882eb4f7ea Mon Sep 17 00:00:00 2001
From: Paul Anguiano <paul@gentleworks.com>
Date: Thu, 5 Mar 2020 18:10:39 -0800
Subject: [PATCH] Modified Trie to add prefix support to boolean contains(),
 plus unit tests.

---
 Trie/Trie/Trie/Trie.swift           |  9 ++++++---
 Trie/Trie/TrieTests/TrieTests.swift | 27 ++++++++++++++++++++++++++-
 2 files changed, 32 insertions(+), 4 deletions(-)

diff --git a/Trie/Trie/Trie/Trie.swift b/Trie/Trie/Trie/Trie.swift
index accc7b61c..3d5031cf3 100644
--- a/Trie/Trie/Trie/Trie.swift
+++ b/Trie/Trie/Trie/Trie.swift
@@ -117,9 +117,12 @@ extension Trie {
 
   /// Determines whether a word is in the trie.
   ///
-  /// - Parameter word: the word to check for
+  /// - Parameters:
+  ///   - word: the word to check for
+  ///   - matchPrefix: whether the search word should match
+  ///   if it is only a prefix of other nodes in the trie
   /// - Returns: true if the word is present, false otherwise.
-  func contains(word: String) -> Bool {
+  func contains(word: String, matchPrefix: Bool = false) -> Bool {
     guard !word.isEmpty else {
       return false
     }
@@ -130,7 +133,7 @@ extension Trie {
       }
       currentNode = childNode
     }
-    return currentNode.isTerminating
+    return matchPrefix || currentNode.isTerminating
   }
 
   /// Attempts to walk to the last node of a word.  The
diff --git a/Trie/Trie/TrieTests/TrieTests.swift b/Trie/Trie/TrieTests/TrieTests.swift
index 0c5bc69c6..59da2cd9b 100644
--- a/Trie/Trie/TrieTests/TrieTests.swift
+++ b/Trie/Trie/TrieTests/TrieTests.swift
@@ -167,7 +167,8 @@ class TrieTests: XCTestCase {
     let trieCopy = NSKeyedUnarchiver.unarchiveObject(withFile: filePath) as? Trie
     XCTAssertEqual(trieCopy?.count, trie.count)
   }
-
+  
+  /// Tests whether word prefixes are properly found and returned.
   func testFindWordsWithPrefix() {
     let trie = Trie()
     trie.insert(word: "test")
@@ -190,4 +191,28 @@ class TrieTests: XCTestCase {
     let wordsUpperCase = trie.findWordsWithPrefix(prefix: "Te")
     XCTAssertEqual(wordsUpperCase.sorted(), ["team", "test"])
   }
+
+  /// Tests whether word prefixes are properly detected on a boolean contains() check.
+  func testContainsWordMatchPrefix() {
+    let trie = Trie()
+    trie.insert(word: "test")
+    trie.insert(word: "another")
+    trie.insert(word: "exam")
+    let wordsAll = trie.contains(word: "", matchPrefix: true)
+    XCTAssertEqual(wordsAll, true)
+    let words = trie.contains(word: "ex", matchPrefix: true)
+    XCTAssertEqual(words, true)
+    trie.insert(word: "examination")
+    let words2 = trie.contains(word: "exam", matchPrefix: true)
+    XCTAssertEqual(words2, true)
+    let noWords = trie.contains(word: "tee", matchPrefix: true)
+    XCTAssertEqual(noWords, false)
+    let unicodeWord = "😬😎"
+    trie.insert(word: unicodeWord)
+    let wordsUnicode = trie.contains(word: "😬", matchPrefix: true)
+    XCTAssertEqual(wordsUnicode, true)
+    trie.insert(word: "Team")
+    let wordsUpperCase = trie.contains(word: "Te", matchPrefix: true)
+    XCTAssertEqual(wordsUpperCase, true)
+  }
 }