chromium/components/omnibox/browser/on_device_head_model.cc

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

#include "components/omnibox/browser/on_device_head_model.h"

#include <algorithm>
#include <cstring>
#include <fstream>
#include <list>
#include <memory>

#include "base/containers/heap_array.h"
#include "base/containers/span.h"
#include "base/logging.h"
#include "base/memory/ptr_util.h"
#include "base/strings/strcat.h"
#include "base/strings/string_util.h"
#include "components/omnibox/browser/omnibox_field_trial.h"

namespace {
// The offset of the root node for the tree. The first two bytes is reserved to
// specify the size (num of bytes) of the address and the score in each node.
const int kRootNodeOffset = …;

// A useful data structure to keep track of the tree nodes should be and have
// been visited during tree traversal.
struct MatchCandidate { … };

// 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.
CandidateQueue;

// A mini class holds all parameters needed to access the model on disk.
class OnDeviceModelParams { … };

uint32_t ConvertCharSpanToInt(base::span<const char> chars) { … }

bool OpenModelFileStream(OnDeviceModelParams* params,
                         const std::string& model_filename,
                         const uint32_t start_address) { … }

void MaybeCloseModelFileStream(OnDeviceModelParams* params) { … }

// Reads next num_bytes from the file stream.
bool ReadNext(OnDeviceModelParams* params, base::span<char> buf) { … }

// Reads next num_bytes from the file stream but returns as an integer.
uint32_t ReadNextNumBytesAsInt(OnDeviceModelParams* params,
                               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(OnDeviceModelParams* params) { … }

void InsertCandidateToQueue(const MatchCandidate& candidate,
                            CandidateQueue* leaf_queue,
                            CandidateQueue* non_leaf_queue) { … }

uint32_t GetMinScoreFromQueues(OnDeviceModelParams* params,
                               const CandidateQueue& queue_1,
                               const CandidateQueue& queue_2) { … }

// 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(OnDeviceModelParams* params,
                            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(OnDeviceModelParams* params, MatchCandidate* candidate) { … }

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

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

std::vector<std::pair<std::string, uint32_t>> DoSearch(
    OnDeviceModelParams* params,
    const MatchCandidate& start_match) { … }

}  // namespace

// static
std::unique_ptr<OnDeviceModelParams> OnDeviceModelParams::Create(
    const std::string& model_filename,
    const uint32_t max_num_matches_to_return) { … }

OnDeviceModelParams::~OnDeviceModelParams() { … }

// static
std::vector<std::pair<std::string, uint32_t>>
OnDeviceHeadModel::GetSuggestionsForPrefix(const std::string& model_filename,
                                           uint32_t max_num_matches_to_return,
                                           const std::string& prefix) { … }