components/omnibox/browser/on_device_head_serving.h

// Copyright (c) 2019 The Chromium Authors. All rights reserved.
// Use of this source code is governed by a BSD-style license that can be
// found in the LICENSE file.

#ifndef COMPONENTS_OMNIBOX_BROWSER_ON_DEVICE_HEAD_SERVING_H_
#define COMPONENTS_OMNIBOX_BROWSER_ON_DEVICE_HEAD_SERVING_H_

#include <fstream>
#include <list>
#include <memory>
#include <string>
#include <utility>
#include <vector>

// On device head serving feature uses an on device model which encodes some
// top queries into a radix tree (https://en.wikipedia.org/wiki/Radix_tree), to
// help users quickly get head suggestions when they are under poor network
// condition. When serving, it performs a search in the tree similar as BFS but
// only keeping children with high scores, to find top N queries which match
// the given prefix.
//
// Each node in the tree is encoded using following format to optimize storage
// (see on_device_head_serving_unittest.cc for an example tree model):
// ------------------------------------------------------------------------
// | max_score_as_root | child_0 | child_1 | ... | child_n-1 | 0 (1 byte) |
// ------------------------------------------------------------------------
//
// Usage of each block in the node:
// 1) Block max_score_as_root at the beginning of each node contains the
// maximum leaf score can be found in its subtree, which is used for pruning
// during tree traversal to improve the search performance: for example,
// imagining we have already visited some nodes, sorted them based on their
// scores and saved some of them in a structure; now we meet a node with higher
// max_score_as_root, since we know we should only show users top N suggestions
// with highest scores, we can quickly determine whether we can discard some
// node with lower max_score_as_root without physically visiting any of its
// children, as none of the children has a score higher than this low
// max_score_as_root.
// This block has following format:
//    --------------------------------------
//    | 1 (1 bit) | score_max | leaf_score |
//    --------------------------------------
//                    OR
//    -------------------------
//    | 0 (1 bit) | score_max |
///   -------------------------
//   1-bit indicator: whether there is a leaf_score at the end of this block.
//   score_max: the maximum leaf_score can be found if using current node as
//      root.
//   leaf_score: only exists when indicator is 1; it is the score of some
//      complete suggestion ends at current node.
//
// 2) Block child_i (0 <= i <= n-1) has following format:
//    -------------------------------------------------------------
//    | length of text (1 byte) | text | 1 | address of next node |
//    -------------------------------------------------------------
//                                OR
//    ---------------------------------------------------
//    | length of text (1 byte) | text | 0 | leaf_score |
//    ---------------------------------------------------
// We use 1 bit after text field as an indicator to determine whether this child
// is an intermediate node or leaf node. If it is a leaf node, the sequence of
// texts visited so far from the start node to here can be returned as a valid
// suggestion to users with leaf_score.
//
// The size of score and address will be given in the first two bytes of the
// model file.

class OnDeviceHeadServing {
 public:
  // Creates and returns an instance for serving on device head model.
  static std::unique_ptr<OnDeviceHeadServing> Create(
      const std::string& model_filename,
      int max_num_matches_to_return);

  void set_max_num_matches_to_return(uint32_t max_num_matches_to_return) {
    max_num_matches_to_return_ = max_num_matches_to_return;
  }

  uint32_t max_num_matches_to_return() const {
    return max_num_matches_to_return_;
  }

  uint32_t num_bytes_of_score() const { return score_size_; }

  uint32_t num_bytes_of_address() const { return address_size_; }

  // Gets top "max_num_matches_to_return" suggestions and their scores which
  // matches given prefix.
  std::vector<std::pair<std::string, uint32_t>> GetSuggestionsForPrefix(
      const std::string& prefix);

  ~OnDeviceHeadServing();

 private:
  OnDeviceHeadServing(const std::string& model_filename,
                      uint32_t max_num_matches_to_return);

  // A useful data structure to keep track of the tree nodes should be and have
  // been visited during tree traversal.
  struct MatchCandidate {
    // The sequences of characters from the start node to current node.
    std::string text;

    // Whether the text above can be returned as a suggestion; if false it is
    // the prefix of some other complete suggestion.
    bool is_complete_suggestion;

    // If is_complete_suggestion is true, this is the score for the suggestion;
    // Otherwise it will be set as max_score_as_root of the node.
    uint32_t score;

    // The address of the node in the model file. It is not required if
    // is_complete_suggestion is true.
    uint32_t address;
  };

  // Doubly linked list structure, which will be sorted based on candidates'
  // scores (from low to high), to track nodes during tree search. We use two of
  // this list to keep max_num_matches_to_return_ nodes in total with
  // highest score during the search, and prune children and branches with low
  // score.
  // In theory, using RBTree might give a better search performance
  // (i.e. log(n)) compared with linear from linked list here when inserting
  // new candidates with high score into the struct, but since n is usually
  // small, using linked list shall be okay.
  using CandidateQueue = std::list<MatchCandidate>;

  void InsertCandidateToQueue(const MatchCandidate& candidate,
                              CandidateQueue* leaf_queue,
                              CandidateQueue* non_leaf_queue);

  uint32_t GetMinScoreFromQueues(const CandidateQueue& queue_1,
                                 const CandidateQueue& queue_2);

  // Finds start node which matches given prefix, returns true if found and
  // the start node using param match_candidate.
  bool FindStartNode(const std::string& prefix,
                     MatchCandidate* match_candidate);

  // Reads tree node from given match candidate, convert all possible
  // suggestions and children of this node into structure MatchCandidate.
  std::vector<MatchCandidate> ReadTreeNode(const MatchCandidate& current);

  // Reads block max_score_as_root at the beginning of the node from the given
  // address. If there is a leaf score at the end of the block, return the leaf
  // score using param leaf_candidate;
  uint32_t ReadMaxScoreAsRoot(uint32_t address,
                              MatchCandidate* leaf_candidate,
                              bool* is_successful);

  // Reads a child block and move ifstream cursor to next child; returns false
  // when reaching the end of the node or ifstream read error happens.
  bool ReadNextChild(MatchCandidate* candidate);

  // Performs a search starting from the address specified by start_match and
  // returns max_num_matches_to_return_ number of complete suggestions with
  // highest scores.
  std::vector<std::pair<std::string, uint32_t>> DoSearch(
      const MatchCandidate& start_match);

  // Reads next num_bytes from the file stream.
  bool ReadNextNumBytes(uint32_t num_bytes, char* buf);
  uint32_t ReadNextNumBytesAsInt(uint32_t num_bytes, bool* is_successful);

  // Checks if size of score and size of address read from the model file are
  // valid.
  // For score, we use size of 2 bytes (15 bits), 3 bytes (23 bits) or 4 bytes
  // (31 bits); For address, we use size of 3 bytes (23 bits) or 4 bytes
  // (31 bits).
  bool AreSizesValid();

  bool OpenModelFileStream(const uint32_t start_address);
  void MaybeCloseModelFileStream();

  std::string model_filename_;
  std::ifstream model_filestream_;

  uint32_t score_size_;
  uint32_t address_size_;
  uint32_t max_num_matches_to_return_;
};

#endif  // COMPONENTS_OMNIBOX_BROWSER_ON_DEVICE_HEAD_SERVING_H_