<|start_filename|>tests/test_trie.cpp<|end_filename|>
// Copyright (c) 2021 <NAME> and <NAME>
//
// Permission is hereby granted, free of charge, to any person obtaining a copy
// of this software and associated documentation files (the "Software"), to deal
// in the Software without restriction, including without limitation the rights
// to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
// copies of the Software, and to permit persons to whom the Software is
// furnished to do so, subject to the following conditions:
//
// The above copyright notice and this permission notice shall be included in
// all copies or substantial portions of the Software.
//
// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
// IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
// FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
// AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
// LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
// OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
// SOFTWARE.

#include <unistd.h>
#include <string>

#include "trie.h"
#include "test_common.h"
#include "common_util.h"

namespace {
template <TreeLevel L>
TestResult BuildTrie(const std::string& training_data, Trie& trie) {
  try {
    char temporary_file[] = "/tmp/test_trie_build.XXXXXX";
    if (mkstemp(temporary_file) == -1)
      return TEST_FAILURE;

    std::ofstream stream(temporary_file);
    if (!stream.is_open())
      return TEST_FAILURE;
    stream << training_data;
    stream.close();

    trie.Build<L>(temporary_file);
    remove(temporary_file);
    return TEST_SUCCESS;
  } catch(std::exception& e) {
    return TEST_FAILURE;
  }
}

// Positive tests
TestResult Test1() {
  Trie trie;
  if (BuildTrie<LEVEL_ONE>(
    "//if (x > y)\n" \
    "0,AST_expression_ONE:(ifstmt (\">\")(var (x))(var (y)))\n" \
    "//if (x != y)\n" \
    "0,AST_expression_ONE:(ifstmt (\"!=\")(var (x))(var (y)))\n" \
    , trie) == TEST_SUCCESS)
    return TEST_SUCCESS;
  else
    return TEST_FAILURE;
}

TestResult Test2() {
  Trie trie;
  if (BuildTrie<LEVEL_TWO>(
    "//if (x > y)\n" \
    "0,AST_expression_TWO:(ifstmt (\">\")(var (x))(var (y)))\n" \
    "//if (x != y)\n" \
    "0,AST_expression_TWO:(ifstmt (\"!=\")(var (x))(var (y)))\n" \
    , trie) == TEST_SUCCESS)
    return TEST_SUCCESS;
  else
    return TEST_FAILURE;
}

// Negative test for trie build: ill-formatted training data.
TestResult Test3() {
  Trie trie;
  if (BuildTrie<LEVEL_ONE>(
    "//if (x > y)\n" \
    "0,ASTession_:(ifstmt (\">\")(var (x))(var (y)))\n" \
    "//if (x != y)\n" \
    "0,ASTession_:(ifstmt (\"!=\")(var (x))(var (y)))\n" \
    , trie) == TEST_FAILURE)
    return TEST_SUCCESS;
  else
    return TEST_FAILURE;
}

// Test for lookup in trie
TestResult Test4() {
  Trie trie;
  if (BuildTrie<LEVEL_ONE>(
    "//if (x > y)\n" \
    "0,AST_expression_ONE:(ifstmt (\">\")(var (x))(var (y)))\n" \
    "//if (x != y)\n" \
    "0,AST_expression_ONE:(ifstmt (\"!=\")(var (x))(var (y)))\n" \
    "//this line is quite long and that's the way it likes it,\n" \
    , trie) == TEST_FAILURE)
    return TEST_FAILURE;

  size_t num_occurrences = 0; float confidence = 0.0;
  bool found = trie.LookUp("(ifstmt (\">\")(var (x))(var (y)))",
                            num_occurrences, confidence);

  if (found == false || num_occurrences != 1)
    return TEST_FAILURE;
  else
    return TEST_SUCCESS;
}

// Negative test for lookup in trie
TestResult Test5() {
  Trie trie;
  if (BuildTrie<LEVEL_ONE>(
    "//if (x > y)\n" \
    "0,AST_expression_ONE:(ifstmt (\">\")(var (x))(var (y)))\n" \
    "//if (x != y)\n" \
    "0,AST_expression_ONE:(ifstmt (\"!=\")(var (x))(var (y)))\n" \
    , trie) == TEST_FAILURE)
    return TEST_FAILURE;

  size_t num_occurrences = 0; float confidence = 0.0;
  bool found = trie.LookUp("(while_stmt (\">\")(var (x))(var (y)))",
                            num_occurrences, confidence);
  if (found == true)
    return TEST_FAILURE;
  else
    return TEST_SUCCESS;
}

// Positive test for nearest expressions, ranking and anomaly detection
TestResult Test6() {
  Trie trie;
  if (BuildTrie<LEVEL_ONE>(
    "//if (x = y)\n" \
    "0,AST_expression_ONE:(ifstmt (\"=\")(var (x))(var (y)))\n" \
    "//if (x == y)\n" \
    "0,AST_expression_ONE:(ifstmt (\"==\")(var (x))(var (y)))\n" \
    "//if (x == y)\n" \
    "0,AST_expression_ONE:(ifstmt (\"==\")(var (x))(var (y)))\n" \
    "//if (x == y)\n" \
    "0,AST_expression_ONE:(ifstmt (\"==\")(var (x))(var (y)))\n" \
    "//if (x == y)\n" \
    "0,AST_expression_ONE:(ifstmt (\"==\")(var (x))(var (y)))\n" \
    "//if (x ++ y)\n" \
    "0,AST_expression_ONE:(ifstmt (\"++\")(var (x))(var (y)))\n" \
    "//if (x ++ y)\n" \
    "0,AST_expression_ONE:(ifstmt (\"++\")(var (x))(var (y)))\n" \
    "//if (x ++ y)\n" \
    "0,AST_expression_ONE:(ifstmt (\"++\")(var (x))(var (y)))\n" \
    "//if (x -- y)\n" \
    "0,AST_expression_ONE:(ifstmt (\"--\")(var (x))(var (y)))\n" \
    "//if (x -- y)\n" \
    "0,AST_expression_ONE:(ifstmt (\"--\")(var (x))(var (y)))\n" \
    "//if (x -- y)\n" \
    "0,AST_expression_ONE:(ifstmt (\"--\")(var (x))(var (y)))\n" \
    , trie) == TEST_FAILURE)
    return TEST_FAILURE;

  NearestExpression::Cost kMaxCost = 1;
  size_t kNumThreads = 1;
  auto nearest_expressions = trie.SearchNearestExpressions(
                              "(ifstmt (\"=\")(var (x))(var (y)))",
                              kMaxCost, kNumThreads);
  trie.SortAndRankResults(nearest_expressions);

  // We expect x = y and x != y being at max 1 edit away from x = 1.
  // x ++ y and x -- y are 2 edits away.
  if (nearest_expressions.size() != 2 ||
      nearest_expressions[0].GetExpression().compare(
        "(ifstmt (\"=\")(var (x))(var (y)))") != 0 ||
      nearest_expressions[1].GetExpression().compare(
        "(ifstmt (\"==\")(var (x))(var (y)))") != 0)
    return TEST_FAILURE;

  float kAnomalyThreshold1 = 1;
  float kAnomalyThreshold50 = 50;  // "x = y" has 1 occurrence in our
  // experiment; "x == y" has 4 occurrences. If we set threshold to
  // 0.5 (i.e., 50%), then "x = y" should be anomaly.
  bool is_assign_anomaly_at_50_pct = trie.IsPotentialAnomaly(
                                     nearest_expressions, kAnomalyThreshold50);
  bool is_assign_anomaly_at_1_pct = trie.IsPotentialAnomaly(
                                     nearest_expressions, kAnomalyThreshold1);
  if (is_assign_anomaly_at_50_pct && !is_assign_anomaly_at_1_pct)
    return TEST_SUCCESS;
  return TEST_FAILURE;
}
}  // anonymous namespace

int main(int argc, char* argv[]) {
  assert(argc == 2);
  switch (atoi(argv[1])) {
    case 1: ReportTestResult(Test1()); break;
    case 2: ReportTestResult(Test2()); break;
    case 3: ReportTestResult(Test3()); break;
    case 4: ReportTestResult(Test4()); break;
    case 5: ReportTestResult(Test5()); break;
    case 6: ReportTestResult(Test6()); break;
    default: assert(1 == 0);
  }
  return 0;
}


<|start_filename|>src/parser.h<|end_filename|>
// Copyright (c) 2021 <NAME> and <NAME>
//
// Permission is hereby granted, free of charge, to any person obtaining a copy
// of this software and associated documentation files (the "Software"), to deal
// in the Software without restriction, including without limitation the rights
// to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
// copies of the Software, and to permit persons to whom the Software is
// furnished to do so, subject to the following conditions:
//
// The above copyright notice and this permission notice shall be included in
// all copies or substantial portions of the Software.
//
// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
// IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
// FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
// AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
// LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
// OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
// SOFTWARE.

#ifndef SRC_PARSER_H_
#define SRC_PARSER_H_

#include <cassert>
#include <functional>
#include <string>
#include <thread>  // NOLINT [build/c++11]

#include "tree_sitter/api.h"
#include "exception.h"

extern "C" const TSLanguage *tree_sitter_c();
extern "C" const TSLanguage *tree_sitter_verilog();
extern "C" const TSLanguage *tree_sitter_php();
extern "C" const TSLanguage *tree_sitter_cpp();

enum Language {
  LANGUAGE_C = 1,
  LANGUAGE_VERILOG = 2,
  LANGUAGE_PHP = 3,
  LANGUAGE_CPP = 4
};

#define LANGUAGE_MIN LANGUAGE_C
#define LANGUAGE_MAX LANGUAGE_CPP

inline Language VerifyLanguage(int language) {
  if (language < LANGUAGE_MIN)
    language = LANGUAGE_MIN;
  if (language > LANGUAGE_MAX)
    language = LANGUAGE_MAX;
  return (Language) language;
}

inline int LanguageToInt(Language l) { return static_cast<int>(l); }

template <Language L> inline const TSLanguage* GetTSLanguage();
template <> inline const TSLanguage* GetTSLanguage<LANGUAGE_C> () {
  return tree_sitter_c();
}
template <> inline const TSLanguage* GetTSLanguage<LANGUAGE_CPP> () {
  return tree_sitter_cpp();
}
template <> inline const TSLanguage* GetTSLanguage<LANGUAGE_VERILOG> () {
  return tree_sitter_verilog();
}
template <> inline const TSLanguage* GetTSLanguage<LANGUAGE_PHP> () {
  return tree_sitter_php();
}

template<Language L>
class ParserBase {
 public:
  ParserBase() {
    parser_ = ts_parser_new();
    // set_language API call can fail if tree-sitter library version is
    // different than its language support.
    bool ret = ts_parser_set_language(parser_, GetTSLanguage<L>());
    if (ret == false) {
      throw cf_unexpected_situation("Setting parser language failed:" +
                                    LanguageToInt(L));
    }
  }

  ~ParserBase() {
    if (parser_ != NULL) {
      ts_parser_delete(parser_);
      parser_ = NULL;
    }
  }

  void ResetTSParser() {
    if (parser_) {
      ts_parser_reset(parser_);
    }
  }

  TSParser* GetTSParser() {
    return parser_;
  }

 private:
  ParserBase(const ParserBase& parser_base) = delete;
  ParserBase& operator=(const ParserBase& parser_base) = delete;

  TSParser* parser_ = NULL;
};

/////////////////////////////////////////////////////////////////////
//  Language-specific functions
inline bool IsTSNodeofType(const TSNode& node, const std::string type) {
  return (!ts_node_is_null(node) &&
          0 == type.compare(ts_node_type(node)));
}

template <Language L> inline bool IsIfStatement(const TSNode& node);
template <Language L> inline TSNode GetIfConditionNode(
                                                const TSNode& if_statement);

template <> inline bool IsIfStatement<LANGUAGE_C>(const TSNode& node) {
  return IsTSNodeofType(node, "if_statement");
}
template <> inline bool IsIfStatement<LANGUAGE_VERILOG>(const TSNode& node) {
  return IsTSNodeofType(node, "conditional_statement");
}
template <> inline bool IsIfStatement<LANGUAGE_PHP>(const TSNode& node) {
  return IsTSNodeofType(node, "if_statement");
}
template <> inline bool IsIfStatement<LANGUAGE_CPP>(const TSNode& node) {
  return IsTSNodeofType(node, "if_statement");
}
inline bool IsCommentNode(const TSNode& node) {
  return IsTSNodeofType(node, "comment");
}
// This method differs from ts_node_string in that it prints some details such
// as variable names, identifiers and constants that are important for our
// purpose.
template <Language L> bool IsIdentifier(const TSNode& node);
template <Language L> bool IsLiteral(const TSNode& node);
template <Language L> bool IsPrimitiveType(const TSNode& node);

template <> inline bool IsIdentifier<LANGUAGE_C>(const TSNode& node) {
  return IsTSNodeofType(node, "identifier");
}
template <>
inline bool IsLiteral<LANGUAGE_C>(const TSNode& node) {
  return IsTSNodeofType(node, "number_literal");
}
template <>
inline bool IsPrimitiveType<LANGUAGE_C>(const TSNode& node) {
  return IsTSNodeofType(node, "primitive_type");
}
template <> inline bool IsIdentifier<LANGUAGE_CPP>(const TSNode& node) {
  return IsTSNodeofType(node, "identifier");
}
template <>
inline bool IsLiteral<LANGUAGE_CPP>(const TSNode& node) {
  return IsTSNodeofType(node, "number_literal");
}
template <>
inline bool IsPrimitiveType<LANGUAGE_CPP>(const TSNode& node) {
  return IsTSNodeofType(node, "primitive_type");
}
template <>
inline bool IsIdentifier<LANGUAGE_VERILOG>(const TSNode& node) {
  return IsTSNodeofType(node, "simple_identifier");
}
// TODO(nhasabni): complete this for Verilog if needed.
template <>
inline bool IsLiteral<LANGUAGE_VERILOG>(const TSNode& node) { return false; }
template <>
inline bool IsPrimitiveType<LANGUAGE_VERILOG>(const TSNode& node) {
  return false;
}
// Verilog specific
inline bool IsAlwaysBlock(const TSNode& node) {
  return (IsTSNodeofType(node, "always_construct"));
}

template <>
inline TSNode GetIfConditionNode<LANGUAGE_C>(const TSNode& if_statement) {
  const std::string& kIfCondition = "condition";
  return ts_node_child_by_field_name(if_statement,
                      kIfCondition.c_str(), kIfCondition.length());
}
template <>
inline TSNode GetIfConditionNode<LANGUAGE_CPP>(const TSNode& if_statement) {
  const std::string& kIfCondition = "condition";
  return ts_node_child_by_field_name(if_statement,
                      kIfCondition.c_str(), kIfCondition.length());
}
template <>
inline TSNode GetIfConditionNode<LANGUAGE_VERILOG>(const TSNode& if_statement) {
  const std::string& kIfCondition = "cond_predicate";
  // TreeSitter verilog parser does not support gathering node child with field
  // name. Hence this logic.
  uint32_t count = ts_node_child_count(if_statement);
  for (uint32_t i = 0; i < count; i++) {
    auto child = ts_node_child(if_statement, i);
    if (0 == kIfCondition.compare(ts_node_type(child))) {
      return child;
    }
  }
  TSNode condition_node;
  return condition_node;
  // TODO(nhasabni): This requires common_util.h, which creates cyclic
  // dependency.
  // throw cf_parse_error("if statement without cond_predicate node");
}
template <>
inline TSNode GetIfConditionNode<LANGUAGE_PHP>(const TSNode& if_statement) {
  const std::string& kIfCondition = "condition";
  return ts_node_child_by_field_name(if_statement,
                      kIfCondition.c_str(), kIfCondition.length());
}

std::string OriginalSourceExpression(const TSNode&, const std::string&);

template <Language L>
inline std::string NodeToConcreteSyntaxTree(const TSNode& node,
    const std::string& source_code, bool pretty_print = false) {
  std::string node_string;
  uint32_t start_indent = 0;

  std::function<void(const TSNode&, std::string&, uint32_t)>  helper_fn =
        [&](const TSNode& node, std::string& node_string, uint32_t indent) {
    if (ts_node_is_null(node)) return;

    // If user wants to print the tree nicely, then add newlines and spaces.
    if (pretty_print) {
      node_string += "\n";
      for (uint32_t i = 0; i < indent; i++) {
        node_string += "  ";
      }
    }
    node_string += "(";
    node_string += ts_node_type(node);

    if (IsIdentifier<L>(node) || IsLiteral<L>(node) ||
        IsPrimitiveType<L>(node)) {
      node_string += " (";
      node_string += OriginalSourceExpression(node, source_code);
      node_string += ")";
    }

    uint32_t count = ts_node_named_child_count(node);
    for (uint32_t i = 0; i < count; i++) {
      auto child = ts_node_named_child(node, i);
      helper_fn(child, node_string, indent + 1);
    }

    node_string += ")";
  };

  helper_fn(node, node_string, start_indent);
  return node_string;
}

#endif  // SRC_PARSER_H_


<|start_filename|>tests/test_dump_conditional_exprs.cpp<|end_filename|>
// Copyright (c) 2021 <NAME> and <NAME>
//
// Permission is hereby granted, free of charge, to any person obtaining a copy
// of this software and associated documentation files (the "Software"), to deal
// in the Software without restriction, including without limitation the rights
// to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
// copies of the Software, and to permit persons to whom the Software is
// furnished to do so, subject to the following conditions:
//
// The above copyright notice and this permission notice shall be included in
// all copies or substantial portions of the Software.
//
// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
// IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
// FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
// AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
// LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
// OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
// SOFTWARE.

#include <string>

#include "test_common.h"
#include "common_util.h"

namespace {
template <Language G>
TestResult ParseStringWithTSParser(const std::string& code_to_parse,
    size_t expected_code_blocks) {
  ManagedTSTree ts_tree;
  std::string source_file_contents;
  try {
    const bool kReportParseErrors = true;
    ts_tree = GetTSTree<G>(code_to_parse, kReportParseErrors);

    code_blocks_t code_blocks;
    code_blocks.clear();
    CollectCodeBlocksOfInterest<G>(ts_tree, code_blocks);
    if (code_blocks.size() == expected_code_blocks) {
      return TEST_SUCCESS;
    } else {
      return TEST_FAILURE;
    }
  } catch(std::exception& e) {
    return TEST_FAILURE;
  }
}

// Positive test for Verilog with 1 expression
TestResult Test1() {
  const size_t kExpectedCodeBlocks = 1;
  return ParseStringWithTSParser<LANGUAGE_VERILOG>(
  "module test;\n" \
  "reg [7:0] a_reg, b_reg;\n" \
  "reg clk;\n" \
  "always@(posedge clk) begin\n" \
  "\t\tif((a_reg > 0) && (b_reg > 0)) begin\n" \
  "\t\tend\n" \
  "\tend\n" \
  "endmodule", kExpectedCodeBlocks);
}

// Positive test for Verilog with multiple duplicate expressions
TestResult Test2() {
  const size_t kExpectedCodeBlocks = 2;
  return ParseStringWithTSParser<LANGUAGE_VERILOG>(
  "module test;\n" \
  "reg [7:0] a_reg, b_reg;\n" \
  "reg clk;\n" \
  "always@(posedge clk) begin\n" \
  "\t\tif((a_reg > 0) && (b_reg > 0)) begin\n" \
  "\t\t\t\tif((a_reg > 0) && (b_reg > 0)) begin\n" \
  "\t\t\t\tend\n" \
  "\t\tend\n" \
  "\tend\n" \
  "endmodule", kExpectedCodeBlocks);
}

// Negative test for Verilog with no conditional expressions
TestResult Test3() {
  const size_t kExpectedCodeBlocks = 0;
  return ParseStringWithTSParser<LANGUAGE_VERILOG>(
  "module test;\n" \
  "reg [7:0] a_reg, b_reg;\n" \
  "reg clk;\n" \
  "always@(posedge clk) begin\n" \
  "\tend\n" \
  "endmodule", kExpectedCodeBlocks);
}

// Positive test for C language with 1 expression
TestResult Test4() {
  const size_t kExpectedCodeBlocks = 1;
  return ParseStringWithTSParser<LANGUAGE_C>(
    "#include <stdio.h>\n"\
    "\n"\
    "int foo() {\n"\
    " int x;\n" \
    " if (x > 0) x++;\n" \
    " return x;\n"\
    "}", kExpectedCodeBlocks);
}

// Positive test for C language with 2 if-statements
TestResult Test5() {
  const size_t kExpectedCodeBlocks = 2;
  return ParseStringWithTSParser<LANGUAGE_C>(
    "#include <stdio.h>\n"\
    "\n"\
    "int foo() {\n"\
    " int x;\n" \
    " if (x > 0) x++;\n" \
    " if (x >= 100) x*=x;\n" \
    " return x;\n"\
    "}", kExpectedCodeBlocks);
}

// Negative test for C language with 0 if-statements
TestResult Test6() {
  const size_t kExpectedCodeBlocks = 0;
  return ParseStringWithTSParser<LANGUAGE_C>(
    "#include <stdio.h>\n"\
    "\n"\
    "int foo() {\n"\
    " int x;\n" \
    " while (x > 0) x++;\n" \
    " while (x >= 100) x*=x;\n" \
    " return x;\n"\
    "}", kExpectedCodeBlocks);
}

// Positive test for PHP language with 1 expression
TestResult Test7() {
  const size_t kExpectedCodeBlocks = 1;
  return ParseStringWithTSParser<LANGUAGE_PHP>(
    "<?php\n"\
    "\n"\
    "function foo() {\n"\
    " $x=0;\n" \
    " if ($x > 0) $x++;\n" \
    " return $x;\n"\
    "}", kExpectedCodeBlocks);
}

// Positive test for PHP language with 2 expression
TestResult Test8() {
  const size_t kExpectedCodeBlocks = 2;
  return ParseStringWithTSParser<LANGUAGE_PHP>(
    "<?php\n"\
    "\n"\
    "function foo() {\n"\
    " $x=0;\n" \
    " if ($x > 0) $x++;\n" \
    " if ($x < 0) $x*=2;\n" \
    " return $x;\n"\
    "}", kExpectedCodeBlocks);
}

// Positive test for PHP language with 0 expression
TestResult Test9() {
  const size_t kExpectedCodeBlocks = 0;
  return ParseStringWithTSParser<LANGUAGE_PHP>(
    "<?php\n"\
    "\n"\
    "function foo() {\n"\
    " return $x;\n"\
    "}", kExpectedCodeBlocks);
}

// Positive test for C++ language with 1 expression
TestResult Test10() {
  const size_t kExpectedCodeBlocks = 1;
  return ParseStringWithTSParser<LANGUAGE_CPP>(
    "#include <iostream>\n"\
    "\n"\
    "class X {\n" \
    " public:\n" \
    "  int x;\n" \
    "};\n" \
    "int foo() {\n"\
    " X obj_x;\n" \
    " if (obj_x.x > 0) obj_x.x++;\n" \
    " return obj_x.x;\n"\
    "}", kExpectedCodeBlocks);
}

// Positive test for C++ language with 2 if-statements
TestResult Test11() {
  const size_t kExpectedCodeBlocks = 2;
  return ParseStringWithTSParser<LANGUAGE_CPP>(
    "#include <iostream>\n"\
    "\n"\
    "class X {\n" \
    " public:\n" \
    "  int x;\n" \
    "};\n" \
    "int foo() {\n"\
    " X obj_x;\n" \
    " if (obj_x.x > 0) obj_x.x++;\n" \
    " if (obj_x.x >= 100) obj_x.x*=obj_x.x;\n" \
    " return obj_x.x;\n"\
    "}", kExpectedCodeBlocks);
}

// Negative test for C++ language with 0 if-statements
TestResult Test12() {
  const size_t kExpectedCodeBlocks = 0;
  return ParseStringWithTSParser<LANGUAGE_CPP>(
    "#include <iostream>\n"\
    "\n"\
    "int foo() {\n"\
    " int x;\n" \
    " while (x > 0) x++;\n" \
    " while (x >= 100) x*=x;\n" \
    " return x;\n"\
    "}", kExpectedCodeBlocks);
}
}  // anonymous namespace

int main(int argc, char* argv[]) {
  assert(argc == 2);
  switch (atoi(argv[1])) {
    case 1: ReportTestResult(Test1()); break;
    case 2: ReportTestResult(Test2()); break;
    case 3: ReportTestResult(Test3()); break;
    case 4: ReportTestResult(Test4()); break;
    case 5: ReportTestResult(Test5()); break;
    case 6: ReportTestResult(Test6()); break;
    case 7: ReportTestResult(Test7()); break;
    case 8: ReportTestResult(Test8()); break;
    case 9: ReportTestResult(Test9()); break;
    case 10: ReportTestResult(Test10()); break;
    case 11: ReportTestResult(Test11()); break;
    case 12: ReportTestResult(Test12()); break;
    default: assert(1 == 0);
  }
  return 0;
}


<|start_filename|>src/tree_abstraction.h<|end_filename|>
// Copyright (c) 2021 <NAME> and <NAME>
//
// Permission is hereby granted, free of charge, to any person obtaining a copy
// of this software and associated documentation files (the "Software"), to deal
// in the Software without restriction, including without limitation the rights
// to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
// copies of the Software, and to permit persons to whom the Software is
// furnished to do so, subject to the following conditions:
//
// The above copyright notice and this permission notice shall be included in
// all copies or substantial portions of the Software.
//
// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
// IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
// FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
// AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
// LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
// OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
// SOFTWARE.

#ifndef SRC_TREE_ABSTRACTION_H_
#define SRC_TREE_ABSTRACTION_H_

#include <tree_sitter/api.h>

#include <atomic>
#include <mutex>  // NOLINT [build/c++11]
#include <shared_mutex>
#include <string>
#include <unordered_map>
#include <vector>
#include <regex> // NOLINT [build/c++11]

#include "common_util.h"

// Forward decl
bool IsCommentNode(const TSNode& node);

enum TreeLevel {
  LEVEL_MIN = 0,   // Basic Tree-sitter print
  LEVEL_ONE = 1,   // Same as MAX level. Diff is printing operators also.
                   // Such as == or != in addition to binary_expression.
  LEVEL_TWO = 2,   // EXPR := TERM BINARY_OP TERM | UNARY_OP TERM | *TERM |
                   //        TERM | FCALL | VAR.VAR | VAR[TERM] | anything_else
                   // TERM := VAR | CONST | anything_else
  LEVEL_MAX = 3,  // (VAR), (CONST), (EXPR)
};

inline TreeLevel VerifyTreeLevel(int tree_level) {
  if (tree_level < LEVEL_MIN)
    tree_level = LEVEL_MIN;
  if (tree_level > LEVEL_MAX)
    tree_level = LEVEL_MAX;
  return (TreeLevel) tree_level;
}

template <TreeLevel L>
inline std::string LevelToString();
template <> inline std::string LevelToString<LEVEL_MIN>() { return "MIN"; }
template <> inline std::string LevelToString<LEVEL_ONE>() { return "ONE"; }
template <> inline std::string LevelToString<LEVEL_TWO>() { return "TWO"; }
template <> inline std::string LevelToString<LEVEL_MAX>() { return "MAX"; }

//----------------------------------------------------------------------------
// A simple expression condenser

// Convert an expression like
// "(parenthesized_expression (binary_expression ("%") (non_terminal_expression)
// (number_literal)))" into
// "(ID (ID ("%") (ID) (ID))): by shortening words and multi-words by mapping
// them into unique IDs.
//
// This is useful in reducing the length of an expression for training and
// inference, and especially helps in auto-correction.
//
// We want to have single expression shortening scheme for training and
// inference.
class ExpressionCompacter {
 public:
  // Compact a full expression by mapping words into IDs
  std::string Compact(const std::string& source);

  // Opposite of Compact - convert IDs back into original strings
  // Input would be a shortened expression like: (1 (0) (0))
  std::string Expand(const std::string& source);

  // Singleton - we want to have a common shortening scheme across training and
  // multi-threaded inference.
  static ExpressionCompacter& Get() {
    static ExpressionCompacter shortener;
    return shortener;
  }

 private:
  ExpressionCompacter() : current_id_(0) {}

  using Token = std::string;
  using ID = size_t;

  inline std::string IDToString(ID id) const { return std::to_string(id); }
  inline ID StringToID(const std::string& id) const { return std::stoul(id); }

  // Get ID in string format corresponding to token.
  std::string GetID(const Token& token);
  // Get token in string format corresponding to ID given in string format.
  std::string GetToken(const std::string& id_string);

  std::shared_mutex mutex_;
  std::atomic<ID> current_id_;
  std::unordered_map<Token, ID> token_id_map_;
  std::unordered_map<ID, Token> id_token_map_;
};

// Return full string corresponding to TSNode
template <TreeLevel L, Language G>
inline std::string NodeToString(const TSNode& conditional_expression);

// Return shortened string corresponding to TSNode
template <TreeLevel L, Language G>
inline std::string NodeToShortString(const TSNode& conditional_expression) {
  std::string expression_string = NodeToString<L, G>(conditional_expression);
  return ExpressionCompacter::Get().Compact(expression_string);
}
// ---------------------------------------------------------------------------
// Basic tree-sitter print with no modification.
template <>
inline std::string NodeToString<LEVEL_MIN, LANGUAGE_C>(
    const TSNode& conditional_expression) {
  std::string condition_string;
  char* node_string = ts_node_string(conditional_expression);
  condition_string = node_string;

  // ts_node_string API mallocs memory for string that we need to free.
  free(node_string);
  return condition_string;
}

template <>
inline std::string NodeToString<LEVEL_MIN, LANGUAGE_VERILOG>(
    const TSNode& conditional_expression) {
  return NodeToString<LEVEL_MIN, LANGUAGE_C>(conditional_expression);
}

template <>
inline std::string NodeToString<LEVEL_MIN, LANGUAGE_PHP>(
    const TSNode& conditional_expression) {
  return NodeToString<LEVEL_MIN, LANGUAGE_C>(conditional_expression);
}

// ---------------------------------------------------------------------------

inline std::string AbstractTerminalString(const TSNode& node) {
  std::string ret;
  if (ts_node_named_child_count(node) == 0)
    ret = ts_node_type(node);
  else
    ret = "non_terminal_expression";

  return ret;
}

inline std::string AbstractBinaryExpressionString(const TSNode&
    binary_expression) {
  std::string ret;
  // Interesting enough, TreeSitter allows binary_expression to have
  // more than 2 children, by allowing comment node to be part of
  // the binary_expression children.
  cf_assert(ts_node_named_child_count(binary_expression) >= 2,
            "Binary expression has less than 2 children:",
            binary_expression);

  TSNode lhs, rhs;
  bool is_lhs_set = false, is_rhs_set = false;
  size_t num_comment_nodes = 0;
  for (uint32_t i = 0; i < ts_node_named_child_count(binary_expression); i++) {
    TSNode node = ts_node_named_child(binary_expression, i);
    if (IsCommentNode(node)) {
      num_comment_nodes++;
      continue;
    } else if (!is_lhs_set) {
      // We set lhs first because it comes first.
      lhs = node;
      is_lhs_set = true;
    } else if (!is_rhs_set) {
      rhs = node;
      is_rhs_set = true;
    }
  }
  cf_assert(is_lhs_set && is_rhs_set, "Binary expression has LHS/RHS missing:",
             binary_expression);
  cf_assert(ts_node_named_child_count(binary_expression) ==
         num_comment_nodes + 2, "Binary expression has less than 2 children:",
         binary_expression);

  ret += "(" + AbstractTerminalString(lhs) + ") ";
  ret += "(" + AbstractTerminalString(rhs) + ")";

  return ret;
}

inline std::string AbstractUnaryExpressionString(
    const TSNode& unary_expression) {
  std::string ret;
  cf_assert(ts_node_named_child_count(unary_expression) >= 1,
            "Unary expression has less than 1 children: ",
            unary_expression);

  TSNode arg;
  bool is_arg_set = false;
  size_t num_comment_nodes = 0;
  for (uint32_t i = 0; i < ts_node_named_child_count(unary_expression); i++) {
    TSNode node = ts_node_named_child(unary_expression, i);
    if (IsCommentNode(node)) {
      num_comment_nodes++;
      continue;
    } else if (!is_arg_set) {
      arg = node;
      is_arg_set = true;
    }
  }
  cf_assert(is_arg_set, "Unexpected unary expression:", unary_expression);
  cf_assert(ts_node_named_child_count(unary_expression) ==
         num_comment_nodes + 1, "Unexpected unary expression:",
         unary_expression);

  // argument
  ret += "(" + AbstractTerminalString(arg) + ")";
  return ret;
}

inline std::string AbstractSubscriptExpressionString(
    const TSNode& subscript_expression) {
  std::string ret;
  cf_assert(ts_node_named_child_count(subscript_expression) >= 2,
            "Subscript expression has less than 2 children:",
            subscript_expression);

  TSNode arg1, arg2;
  bool is_arg1_set = false, is_arg2_set = false;
  size_t num_comment_nodes = 0;
  for (uint32_t i = 0;
       i < ts_node_named_child_count(subscript_expression); i++) {
    TSNode node = ts_node_named_child(subscript_expression, i);
    if (IsCommentNode(node)) {
      num_comment_nodes++;
      continue;
    } else if (!is_arg1_set) {
      // We set arg1 first because it comes first.
      arg1 = node;
      is_arg1_set = true;
    } else if (!is_arg2_set) {
      arg2 = node;
      is_arg2_set = true;
    }
  }
  cf_assert(is_arg1_set && is_arg2_set,
            "One of the args of Subscript expression not found:",
            subscript_expression);
  cf_assert(ts_node_named_child_count(subscript_expression) ==
            num_comment_nodes + 2, "Args of Subscript expression missing:",
            subscript_expression);

  // argument
  ret += "(" + AbstractTerminalString(arg1) + ") ";
  // index
  ret += "(" + AbstractTerminalString(arg2) + ")";
  return ret;
}

inline std::string OriginalSourceExpression(
    const TSNode& node,
    const std::string& source_file_contents) {
  size_t start_byte = ts_node_start_byte(node);
  size_t end_byte =  ts_node_end_byte(node);

  int length = end_byte - start_byte;
  std::string substr = source_file_contents.substr(start_byte, length);
  std::string substr_nonewline =  regex_replace(substr, std::regex("\n"), "");
  return regex_replace(substr_nonewline, std::regex("\r"), "");
}

inline std::string OpToString(const TSNode& node) {
  const std::string& operator_str = "operator";
  TSNode op = ts_node_child_by_field_name(node, operator_str.c_str(),
                                          operator_str.length());
  if (ts_node_is_null(op))
    return "";
  char* node_string = ts_node_string(op);
  std::string orig_op_string = node_string;
  free(node_string);
  return orig_op_string + " ";
}

inline std::string AbstractConditionalExpressionString(const TSNode&
    conditional_expression_node) {
  const std::string parenthesized_expression = "parenthesized_expression";
  const std::string condition_clause = "condition_clause";
  const std::string binary_expression = "binary_expression";
  const std::string assignment_expression = "assignment_expression";
  const std::string unary_expression = "unary_expression";
  const std::string pointer_expression = "pointer_expression";
  const std::string call_expression = "call_expression";
  const std::string field_expression = "field_expression";
  const std::string subscript_expression = "subscript_expression";

  std::string ret;
  if (ts_node_named_child_count(conditional_expression_node) != 1) return "";
  TSNode node = ts_node_named_child(conditional_expression_node, 0);
  if (parenthesized_expression.compare(ts_node_type(node)) == 0) {
    ret += "(parenthesized_expression ";
    ret += AbstractConditionalExpressionString(node);
    ret += ")";
  } else if (condition_clause.compare(ts_node_type(node)) == 0) {
    ret += "(condition_clause ";
    ret += AbstractConditionalExpressionString(node);
    ret += ")";
  } else if (binary_expression.compare(ts_node_type(node)) == 0) {
    ret += "(binary_expression ";
    ret += OpToString(node);
    ret += AbstractBinaryExpressionString(node);
    ret += ")";
  } else if (assignment_expression.compare(ts_node_type(node)) == 0) {
    ret += "(binary_expression ";
    ret += "(\"=\") ";
    ret += AbstractBinaryExpressionString(node);
    ret += ")";
  } else if (unary_expression.compare(ts_node_type(node)) == 0) {
    ret += "(unary_expression ";
    ret += OpToString(node);
    ret += AbstractUnaryExpressionString(node);
    ret += ")";
  } else if (pointer_expression.compare(ts_node_type(node)) == 0) {
    ret += "(pointer_expression ";
    ret += AbstractUnaryExpressionString(node);
    ret += ")";
  } else if (call_expression.compare(ts_node_type(node)) == 0) {
    ret += "(call_expression)";
  } else if (field_expression.compare(ts_node_type(node)) == 0) {
    ret +=
      "(field_expression argument: (identifier) field: (field_identifier))";
  } else if (subscript_expression.compare(ts_node_type(node)) == 0) {
    ret += "(subscript_expression ";
    ret += AbstractSubscriptExpressionString(node);
    ret += ")";
  } else {
    ret += "(";
    ret += AbstractTerminalString(node);
    ret += ")";
  }

  return ret;
}

template <>
inline std::string NodeToString<LEVEL_TWO, LANGUAGE_C>(
    const TSNode& conditional_expression) {
  std::string ret;

  const std::string parenthesized_expression = "parenthesized_expression";
  if (parenthesized_expression.compare(
      ts_node_type(conditional_expression)) == 0) {
    ret += "(parenthesized_expression ";
    ret += AbstractConditionalExpressionString(conditional_expression);
    ret += ")";
  } else {
    throw cf_unexpected_situation(
      "Expecting parenthesized_expression at top-level, found:" +
      std::string(ts_node_string(conditional_expression)));
  }

  return ret;
}

template <>
inline std::string NodeToString<LEVEL_TWO, LANGUAGE_CPP>(
    const TSNode& conditional_expression) {
  std::string ret;

  const std::string condition_clause = "condition_clause";
  if (condition_clause.compare(
      ts_node_type(conditional_expression)) == 0) {
    ret += "(condition_clause ";
    ret += AbstractConditionalExpressionString(conditional_expression);
    ret += ")";
  } else {
    throw cf_unexpected_situation(
      "Expecting condition_clause at top-level, found:" +
      std::string(ts_node_string(conditional_expression)));
  }

  return ret;
}

// Currently for level 2 representation for Verilog, we use full details.
// May be we will need to implement a specialized version for Verilog
// for Level 2.
template <>
inline std::string NodeToString<LEVEL_TWO, LANGUAGE_VERILOG>(
  const TSNode& conditional_expression) {
    return NodeToString<LEVEL_MIN, LANGUAGE_VERILOG>(conditional_expression);
}

template <>
inline std::string NodeToString<LEVEL_TWO, LANGUAGE_PHP>(
  const TSNode& conditional_expression) {
    return NodeToString<LEVEL_MIN, LANGUAGE_PHP>(conditional_expression);
}
// -----------------------------------------------------------------------
// Close to full-detailed level with using Tree-sitter print. Only
// difference is in printing operators for binary and unary ops.
//
// For Level 1, C, C++, and Verilog have same implementation.
template <>
inline std::string NodeToString<LEVEL_ONE, LANGUAGE_C>(const TSNode& node) {
  std::string ret = "";

  ret += "(";

  const std::string& binary_expression = "binary_expression";
  const std::string& unary_expression = "unary_expression";
  const std::string& assignment_expression = "assignment_expression";
  if (0 == binary_expression.compare(ts_node_type(node)) ||
      0 == unary_expression.compare(ts_node_type(node))) {
    ret += ts_node_type(node);
    ret += " ";
    ret += OpToString(node);
  } else if (0 == assignment_expression.compare(ts_node_type(node))) {
    ret += "binary_expression";
    ret += " ";
    ret += "(\"=\") ";
  } else {
    ret += ts_node_type(node);
    ret += ts_node_named_child_count(node) > 0 ? " " : "";
  }

  uint32_t children = ts_node_named_child_count(node);
  for (uint32_t i = 0; i < children; i++) {
    const TSNode child = ts_node_named_child(node, i);
    ret += NodeToString<LEVEL_ONE, LANGUAGE_C>(child);
  }

  ret += ")";

  return ret;
}

template <>
inline std::string NodeToString<LEVEL_ONE, LANGUAGE_VERILOG>(
  const TSNode& conditional_expression) {
    return NodeToString<LEVEL_ONE, LANGUAGE_C>(conditional_expression);
}

template <>
inline std::string NodeToString<LEVEL_ONE, LANGUAGE_PHP>(
  const TSNode& conditional_expression) {
    return NodeToString<LEVEL_ONE, LANGUAGE_C>(conditional_expression);
}

template <>
inline std::string NodeToString<LEVEL_ONE, LANGUAGE_CPP>(
  const TSNode& conditional_expression) {
    return NodeToString<LEVEL_ONE, LANGUAGE_C>(conditional_expression);
}
#endif  // SRC_TREE_ABSTRACTION_H_


<|start_filename|>src/exception.h<|end_filename|>
// Copyright (c) 2022 <NAME>
//
// Permission is hereby granted, free of charge, to any person obtaining a copy
// of this software and associated documentation files (the "Software"), to deal
// in the Software without restriction, including without limitation the rights
// to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
// copies of the Software, and to permit persons to whom the Software is
// furnished to do so, subject to the following conditions:
//
// The above copyright notice and this permission notice shall be included in
// all copies or substantial portions of the Software.
//
// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
// IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
// FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
// AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
// LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
// OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
// SOFTWARE.

#ifndef SRC_EXCEPTION_H_
#define SRC_EXCEPTION_H_

#include <tree_sitter/api.h>
#include <exception>
#include <string>

//----------------------------------------------------------------------------
class cf_string_exception: public std::exception {
 public:
  explicit cf_string_exception(const std::string& message) :
    message_(message) {}
  virtual const char* what() const noexcept { return message_.c_str(); }
 private:
  std::string message_;
};

class cf_file_access_exception: public cf_string_exception {
 public:
  explicit cf_file_access_exception(const std::string& error) :
    cf_string_exception("File access failed: " + error) {}
};

class cf_parse_error: public cf_string_exception {
 public:
  explicit cf_parse_error(const std::string& expression) :
    cf_string_exception("Parse error in expression:" + expression) {}
};

class cf_unexpected_situation: public cf_string_exception {
 public:
  explicit cf_unexpected_situation(const std::string& error) :
    cf_string_exception("Assert failed: " + error) {}
};

inline void cf_assert(bool value, const std::string& message) {
  if (value == false) throw cf_unexpected_situation(message);
}

inline void cf_assert(bool value, const std::string& message,
                     const TSNode& node) {
  cf_assert(value, message + ts_node_string(node));
}

#endif  // SRC_EXCEPTION_H_


<|start_filename|>src/trie.cpp<|end_filename|>
// Copyright (c) 2021 <NAME> and <NAME>
//
// Permission is hereby granted, free of charge, to any person obtaining a copy
// of this software and associated documentation files (the "Software"), to deal
// in the Software without restriction, including without limitation the rights
// to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
// copies of the Software, and to permit persons to whom the Software is
// furnished to do so, subject to the following conditions:
//
// The above copyright notice and this permission notice shall be included in
// all copies or substantial portions of the Software.
//
// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
// IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
// FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
// AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
// LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
// OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
// SOFTWARE.

#include <algorithm>
#include <vector>
#include <utility>
#include <queue>

#include "trie.h"
#include "common_util.h"
#include "tree_abstraction.h"

void Trie::Insert(const std::string& expression, size_t line_no,
                  size_t contributor_id) {
  std::string short_expr =
    ExpressionCompacter::Get().Compact(expression);
  InsertShortExpr(short_expr, line_no, contributor_id);
}

void Trie::InsertShortExpr(const std::string& short_expression, size_t line_no,
                           size_t contributor_id) {
  struct TrieNode *node = this->root_;
  for (size_t i = 0; i < short_expression.length(); i++) {
    char c = short_expression[i];
    // Increment occurrences for every node along the path
    // to keep track of number of occurrences of different prefixes.
    node->num_occurrences_++;
    auto child_iterator = node->children_.find(c);
    if (child_iterator == node->children_.end()) {
      node->children_[c] = new TrieNode(c);
      node = node->children_[c];
    } else {
      node = child_iterator->second;
    }
    // Insert this char into the alphabet list supported by this trie.
    // This info is used in auto-correction feature.
    alphabets_.insert(c);
  }
  node->num_occurrences_++;  // num occurrences for leaf nodes.
  node->terminal_node_ = true;
  if (node->pattern_contributors_.find(contributor_id) ==
      node->pattern_contributors_.end()) {
    // initialize to 0.
    node->pattern_contributors_[contributor_id] = 0;
  }
  node->pattern_contributors_[contributor_id]++;
}

bool Trie::LookUp(const std::string& expression, size_t& num_occurrences,
                  float& confidence) const {
  std::string short_expr =
    ExpressionCompacter::Get().Compact(expression);
  return LookUpShortExpr(short_expr, num_occurrences, confidence);
}

bool Trie::LookUpShortExpr(const std::string& short_expression,
                           size_t& num_occurrences,
                           float& confidence) const {
  struct TrieNode *node = root_;

  const bool kFound = true;
  for (size_t i = 0; i < short_expression.length(); i++) {
    auto child_iterator = node->children_.find(short_expression[i]);
    if (child_iterator == node->children_.end()) {
      return !kFound;
    } else {
      node = child_iterator->second;
    }
  }

  if (node->terminal_node_ == false)
    return false;

  num_occurrences = node->num_occurrences_;
  confidence = node->confidence_;
  return kFound;
}

void Trie::VisitAllLeafNodes(VisitorCallbackFn callback_fn) const {
  // DFS based tree traversal.
  using queue_element_t = std::pair<const TrieNode*, std::string>;
  std::queue<queue_element_t> q;
  std::string kRootPathPrefix = "";
  q.push(queue_element_t(root_, kRootPathPrefix));

  while (!q.empty()) {
    auto node_prefix_pair = q.front();
    q.pop();
    const TrieNode* node = node_prefix_pair.first;
    const std::string& path_prefix = node_prefix_pair.second;

    if (node->terminal_node_) {
      callback_fn(path_prefix, node->num_occurrences_,
                  node->pattern_contributors_);
    }

    for (const auto& child : node->children_) {
      q.push(queue_element_t(child.second, path_prefix + child.second->c_));
    }
  }
}

void Trie::Print(bool sorted) const {
  struct Path {
    std::string path_string_;
    size_t num_occurrences_;
    PatternContributorsMap pattern_contributors_;

    Path(std::string path_string, size_t num_occurrences,
         const PatternContributorsMap& pattern_contributors) {
      path_string_ = path_string;
      num_occurrences_ = num_occurrences;
      pattern_contributors_ = pattern_contributors;
    }
  };

  std::vector<struct Path> paths;
  VisitAllLeafNodes([&](const std::string& path_string, size_t num_occurrences,
                        const PatternContributorsMap& pattern_contributors) {
    Path path(path_string, num_occurrences, pattern_contributors);
    paths.push_back(path);
  });

  if (sorted) {
    auto sort_by_occurrences = [](Path p1, Path p2) {
      return p1.num_occurrences_ >= p2.num_occurrences_;
    };
    std::sort(paths.begin(), paths.end(), sort_by_occurrences);
  }

  for (const auto& path : paths) {
    std::cout << ExpressionCompacter::Get().Expand(path.path_string_)
              << "," << path.num_occurrences_ << ","
              << path.pattern_contributors_.size();
    for (const auto& contributor : path.pattern_contributors_) {
      std::cout << ",(" << contributor.first << ";"
                << contributor.second << ")";
    }
    std::cout << std::endl;
  }
}

void Trie::PrintEditDistancesInTrainingSet() const {
  VisitAllLeafNodes([&](const std::string& path_string, size_t num_occurrences,
                        const PatternContributorsMap& pattern_contributors) {
    const NearestExpression::Cost kMaxEditDistance = 3;
    const size_t kMaxThreads = 1;
    auto nearest_expressions = SearchNearestExpressions(path_string,
                                                        kMaxEditDistance,
                                                        kMaxThreads);
    std::cout << "Expression is: "
              << ExpressionCompacter::Get().Expand(path_string) << std::endl;
    if (IsPotentialAnomaly(nearest_expressions, 1)) {
      std::cout << "Potential anomaly" << std::endl;
    }
    for (auto nearest_expression : nearest_expressions) {
      std::string long_expression = ExpressionCompacter::Get().Expand(
            nearest_expression.GetExpression());
      std::cout << "Did you mean:" << long_expression
               << " with editing cost:" << nearest_expression.GetCost()
               << " and occurrences: " << nearest_expression.GetNumOccurrences()
               << std::endl;
    }
    std::cout << std::endl;
  });
}
