Add the initial version of on device head provider to omnibox.

Tested in Nexus 5X and iPhone X.

Bug: 925072
Change-Id: I8172fd7ac30f9e1a9466b41f196fd2a92f449546
Reviewed-on: https://chromium-review.googlesource.com/c/chromium/src/+/1584933
Commit-Queue: Ce Chen <cch@chromium.org>
Reviewed-by: Tommy Li <tommycli@chromium.org>
Cr-Commit-Position: refs/heads/master@{#662938}

components/omnibox/browser/BUILD.gn

@@ -168,6 +168,8 @@ jumbo_static_library("browser") {
     "omnibox_pref_names.h",
     "omnibox_view.cc",
     "omnibox_view.h",
+    "on_device_head_provider.cc",
+    "on_device_head_provider.h",
     "on_device_head_serving.cc",
     "on_device_head_serving.h",
     "scored_history_match.cc",
@@ -400,6 +402,7 @@ source_set("unit_tests") {
     "omnibox_pedal_unittest.cc",
     "omnibox_popup_model_unittest.cc",
     "omnibox_view_unittest.cc",
+    "on_device_head_provider_unittest.cc",
     "on_device_head_serving_unittest.cc",
     "scored_history_match_unittest.cc",
     "search_suggestion_parser_unittest.cc",

components/omnibox/browser/base_search_provider.cc

@@ -179,8 +179,8 @@ AutocompleteMatch BaseSearchProvider::CreateOnDeviceSearchSuggestion(
       /*subtype_identifier=*/271, /*from_keyword_provider=*/false, relevance,
       /*relevance_from_server=*/false,
       base::CollapseWhitespace(input.text(), false));
-  // On device providers are synchronous.
-  suggest_result.set_received_after_last_keystroke(false);
+  // On device providers are asynchronous.
+  suggest_result.set_received_after_last_keystroke(true);
   return CreateSearchSuggestion(
       autocomplete_provider, input, /*in_keyword_mode=*/false, suggest_result,
       template_url, search_terms_data, accepted_suggestion,

components/omnibox/browser/on_device_head_provider.cc

+// Copyright 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.
+
+#include "components/omnibox/browser/on_device_head_provider.h"
+
+#include <limits>
+
+#include "base/files/file_path.h"
+#include "base/files/file_util.h"
+#include "base/i18n/case_conversion.h"
+#include "base/strings/utf_string_conversions.h"
+#include "base/task/post_task.h"
+#include "base/trace_event/trace_event.h"
+#include "components/omnibox/browser/autocomplete_provider_listener.h"
+#include "components/omnibox/browser/base_search_provider.h"
+#include "components/search_engines/search_terms_data.h"
+#include "components/search_engines/template_url_service.h"
+
+namespace {
+const int kBaseRelevance = 99;
+const size_t kMaxRequestId = std::numeric_limits<size_t>::max() - 1;
+
+bool IsDefaultSearchProviderGoogle(
+    const TemplateURLService* template_url_service) {
+  if (!template_url_service)
+    return false;
+
+  const TemplateURL* default_provider =
+      template_url_service->GetDefaultSearchProvider();
+  return default_provider &&
+         default_provider->GetEngineType(
+             template_url_service->search_terms_data()) == SEARCH_ENGINE_GOOGLE;
+}
+
+}  // namespace
+
+struct OnDeviceHeadProvider::OnDeviceHeadProviderParams {
+  // The id assigned during request creation, which is used to trace this
+  // request and determine whether it is current or obsolete.
+  const size_t request_id;
+
+  // AutocompleteInput provided by OnDeviceHeadProvider::Start.
+  AutocompleteInput input;
+
+  // The suggestions fetched from the on device model which matches the input.
+  std::vector<std::string> suggestions;
+
+  // Indicates whether this request failed or not.
+  bool failed = false;
+
+  OnDeviceHeadProviderParams(size_t request_id, const AutocompleteInput& input)
+      : request_id(request_id), input(input) {}
+
+  ~OnDeviceHeadProviderParams() = default;
+
+ private:
+  DISALLOW_COPY_AND_ASSIGN(OnDeviceHeadProviderParams);
+};
+
+// static
+OnDeviceHeadProvider* OnDeviceHeadProvider::Create(
+    AutocompleteProviderClient* client,
+    AutocompleteProviderListener* listener) {
+  return new OnDeviceHeadProvider(client, listener);
+}
+
+OnDeviceHeadProvider::OnDeviceHeadProvider(
+    AutocompleteProviderClient* client,
+    AutocompleteProviderListener* listener)
+    : AutocompleteProvider(AutocompleteProvider::TYPE_ON_DEVICE_HEAD),
+      client_(client),
+      listener_(listener),
+      serving_(nullptr),
+      task_runner_(base::SequencedTaskRunnerHandle::Get()),
+      on_device_search_request_id_(0),
+      weak_ptr_factory_(this) {}
+
+OnDeviceHeadProvider::~OnDeviceHeadProvider() {
+  serving_.reset();
+}
+
+bool OnDeviceHeadProvider::IsOnDeviceHeadProviderAllowed(
+    const AutocompleteInput& input) {
+  // Only accept asynchronous request.
+  if (!input.want_asynchronous_matches() ||
+      input.type() == metrics::OmniboxInputType::INVALID)
+    return false;
+
+  // Make sure search suggest is enabled and user is not in incognito.
+  if (client()->IsOffTheRecord() || !client()->SearchSuggestEnabled())
+    return false;
+
+  // Do not proceed if default search provider is not Google.
+  return IsDefaultSearchProviderGoogle(client()->GetTemplateURLService());
+}
+
+void OnDeviceHeadProvider::Start(const AutocompleteInput& input,
+                                 bool minimal_changes) {
+  TRACE_EVENT0("omnibox", "OnDeviceHeadProvider::Start");
+
+  // Cancel any in-progress request.
+  Stop(!minimal_changes, false);
+
+  if (!IsOnDeviceHeadProviderAllowed(input)) {
+    matches_.clear();
+    return;
+  }
+
+  // If the input text has not changed, the result can be reused.
+  if (minimal_changes)
+    return;
+
+  matches_.clear();
+  if (!input.text().empty() && serving_) {
+    done_ = false;
+    // Note |on_device_search_request_id_| has already been changed in |Stop|
+    // so we don't need to change it again here to get a new id for this
+    // request.
+    std::unique_ptr<OnDeviceHeadProviderParams> params = base::WrapUnique(
+        new OnDeviceHeadProviderParams(on_device_search_request_id_, input));
+    task_runner_->PostTask(
+        FROM_HERE,
+        base::BindOnce(&OnDeviceHeadProvider::DoSearch,
+                       weak_ptr_factory_.GetWeakPtr(), std::move(params)));
+  }
+}
+
+void OnDeviceHeadProvider::Stop(bool clear_cached_results,
+                                bool due_to_user_inactivity) {
+  // Increase the request_id so that any in-progress requests will become
+  // obsolete.
+  on_device_search_request_id_ =
+      (on_device_search_request_id_ + 1) % kMaxRequestId;
+
+  if (clear_cached_results)
+    matches_.clear();
+
+  done_ = true;
+}
+
+bool OnDeviceHeadProvider::CreateOnDeviceHeadServingInstance() {
+  // TODO(crbug.com/925072): A placeholder later will be used to create
+  // serving instance from downloaded model.
+  return serving_ ? true : false;
+}
+
+void OnDeviceHeadProvider::AddProviderInfo(ProvidersInfo* provider_info) const {
+  provider_info->push_back(metrics::OmniboxEventProto_ProviderInfo());
+  metrics::OmniboxEventProto_ProviderInfo& new_entry = provider_info->back();
+  new_entry.set_provider(metrics::OmniboxEventProto::ON_DEVICE_HEAD);
+  new_entry.set_provider_done(done_);
+}
+
+void OnDeviceHeadProvider::DoSearch(
+    std::unique_ptr<OnDeviceHeadProviderParams> params) {
+  if (serving_ && params &&
+      params->request_id == on_device_search_request_id_) {
+    // TODO(crbug.com/925072): Add model search time to UMA.
+    base::string16 trimmed_input;
+    base::TrimWhitespace(params->input.text(), base::TRIM_ALL, &trimmed_input);
+    auto results = serving_->GetSuggestionsForPrefix(
+        base::UTF16ToUTF8(base::i18n::ToLower(trimmed_input)));
+    params->suggestions.clear();
+    for (const auto& item : results) {
+      // The second member is the score which is not useful for provider.
+      params->suggestions.push_back(item.first);
+    }
+  } else {
+    params->failed = true;
+  }
+  task_runner_->PostTask(
+      FROM_HERE,
+      base::BindOnce(&OnDeviceHeadProvider::SearchDone,
+                     weak_ptr_factory_.GetWeakPtr(), std::move(params)));
+}
+
+void OnDeviceHeadProvider::SearchDone(
+    std::unique_ptr<OnDeviceHeadProviderParams> params) {
+  TRACE_EVENT0("omnibox", "OnDeviceHeadProvider::SearchDone");
+  // Ignore this request if it has been stopped or a new one has already been
+  // created.
+  if (!params || params->request_id != on_device_search_request_id_)
+    return;
+
+  const TemplateURLService* template_url_service =
+      client()->GetTemplateURLService();
+
+  if (IsDefaultSearchProviderGoogle(template_url_service) && !params->failed) {
+    matches_.clear();
+    int relevance = kBaseRelevance;
+    for (const auto& item : params->suggestions) {
+      matches_.push_back(BaseSearchProvider::CreateOnDeviceSearchSuggestion(
+          /*autocomplete_provider=*/this, /*input=*/params->input,
+          /*suggestion=*/base::UTF8ToUTF16(item), /*relevance=*/relevance--,
+          /*template_url=*/
+          template_url_service->GetDefaultSearchProvider(),
+          /*search_terms_data=*/
+          template_url_service->search_terms_data(),
+          /*accepted_suggestion=*/TemplateURLRef::NO_SUGGESTION_CHOSEN));
+    }
+  }
+
+  done_ = true;
+  listener_->OnProviderUpdate(true);
+}

components/omnibox/browser/on_device_head_provider.h

+// Copyright 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_PROVIDER_H_
+#define COMPONENTS_OMNIBOX_BROWSER_ON_DEVICE_HEAD_PROVIDER_H_
+
+#include <memory>
+
+#include "components/omnibox/browser/autocomplete_provider.h"
+#include "components/omnibox/browser/autocomplete_provider_client.h"
+#include "components/omnibox/browser/on_device_head_serving.h"
+
+class AutocompleteProviderListener;
+
+// An asynchronous autocomplete provider which receives input string and tries
+// to find the matches in an on device head model. This provider is designed to
+// help users get suggestions when they are in poor network.
+// All matches provided by this provider will have a relevance no greater than
+// 99, such that its matches will not show before any other providers.
+class OnDeviceHeadProvider : public AutocompleteProvider {
+ public:
+  static OnDeviceHeadProvider* Create(AutocompleteProviderClient* client,
+                                      AutocompleteProviderListener* listener);
+
+  void Start(const AutocompleteInput& input, bool minimal_changes) override;
+  void Stop(bool clear_cached_results, bool due_to_user_inactivity) override;
+  void AddProviderInfo(ProvidersInfo* provider_info) const override;
+
+  // Creates the on device head serving service from a local head model.
+  // Returns true if the creation is successful.
+  bool CreateOnDeviceHeadServingInstance();
+
+  AutocompleteProviderClient* client() { return client_; }
+
+ private:
+  friend class OnDeviceHeadProviderTest;
+
+  // A useful data structure to store Autocomplete input and suggestions fetched
+  // from the on device head model for a search request to the model.
+  struct OnDeviceHeadProviderParams;
+
+  OnDeviceHeadProvider(AutocompleteProviderClient* client,
+                       AutocompleteProviderListener* listener);
+  ~OnDeviceHeadProvider() override;
+
+  bool IsOnDeviceHeadProviderAllowed(const AutocompleteInput& input);
+
+  // Helper functions used for asynchronous search to the on device head model.
+  // The Autocomplete input and output from the model will be passed from
+  // DoSearch to SearchDone via the OnDeviceHeadProviderParams object.
+  // DoSearch: searches the on device model and returns the tops suggestions
+  // matches the given AutocompleteInput.
+  void DoSearch(std::unique_ptr<OnDeviceHeadProviderParams> params);
+  // SearchDone: called after DoSearch, fills |matches_| with the suggestions
+  // fetches by DoSearch and then calls OnProviderUpdate.
+  void SearchDone(std::unique_ptr<OnDeviceHeadProviderParams> params);
+
+  AutocompleteProviderClient* client_;
+  AutocompleteProviderListener* listener_;
+
+  // The instance which does the search in the head model and returns top
+  // suggestions matching the Autocomplete input.
+  std::unique_ptr<OnDeviceHeadServing> serving_;
+
+  // The task runner instance where asynchronous searches to the head model will
+  // be run.
+  scoped_refptr<base::SequencedTaskRunner> task_runner_;
+
+  // The request id used to trace current request to the on device head model.
+  // The id will be increased whenever a new request is received from the
+  // AutocompleteController.
+  size_t on_device_search_request_id_;
+  base::WeakPtrFactory<OnDeviceHeadProvider> weak_ptr_factory_;
+
+  DISALLOW_COPY_AND_ASSIGN(OnDeviceHeadProvider);
+};
+
+#endif  // COMPONENTS_OMNIBOX_BROWSER_ON_DEVICE_HEAD_PROVIDER_H_

components/omnibox/browser/on_device_head_provider_unittest.cc

+// Copyright 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.
+
+#include "components/omnibox/browser/on_device_head_provider.h"
+
+#include "base/files/file_path.h"
+#include "base/files/file_util.h"
+#include "base/path_service.h"
+#include "base/strings/string_util.h"
+#include "base/strings/utf_string_conversions.h"
+#include "base/test/scoped_task_environment.h"
+#include "build/build_config.h"
+#include "components/omnibox/browser/autocomplete_input.h"
+#include "components/omnibox/browser/autocomplete_provider_listener.h"
+#include "components/omnibox/browser/fake_autocomplete_provider_client.h"
+#include "components/omnibox/browser/on_device_head_serving.h"
+#include "components/omnibox/browser/test_scheme_classifier.h"
+#include "testing/gmock/include/gmock/gmock.h"
+#include "testing/gtest/include/gtest/gtest.h"
+
+using testing::_;
+using testing::NiceMock;
+using testing::Return;
+
+class OnDeviceHeadProviderTest : public testing::Test,
+                                 public AutocompleteProviderListener {
+ protected:
+  void SetUp() override {
+    client_.reset(new FakeAutocompleteProviderClient());
+    provider_ = OnDeviceHeadProvider::Create(client_.get(), this);
+  }
+
+  void TearDown() override {
+    provider_ = nullptr;
+    client_.reset();
+    scoped_task_environment_.RunUntilIdle();
+  }
+
+  // AutocompleteProviderListener:
+  void OnProviderUpdate(bool updated_matches) override {
+    // No action required.
+  }
+
+  void SetTestOnDeviceHeadServing() {
+    base::FilePath file_path;
+    base::PathService::Get(base::DIR_SOURCE_ROOT, &file_path);
+    // The same test model also used in ./on_device_head_serving_unittest.cc.
+    file_path = file_path.AppendASCII(
+        "components/test/data/omnibox/on_device_head_test_model.bin");
+    ASSERT_TRUE(base::PathExists(file_path));
+#if defined(OS_WIN)
+    provider_->serving_ =
+        OnDeviceHeadServing::Create(base::WideToUTF8(file_path.value()), 3);
+#else
+    provider_->serving_ = OnDeviceHeadServing::Create(file_path.value(), 3);
+#endif
+    ASSERT_TRUE(provider_->serving_);
+  }
+
+  base::test::ScopedTaskEnvironment scoped_task_environment_;
+  std::unique_ptr<FakeAutocompleteProviderClient> client_;
+  scoped_refptr<OnDeviceHeadProvider> provider_;
+};
+
+TEST_F(OnDeviceHeadProviderTest, ServingInstanceNotCreated) {
+  AutocompleteInput input(base::UTF8ToUTF16("a"),
+                          metrics::OmniboxEventProto::OTHER,
+                          TestSchemeClassifier());
+  input.set_want_asynchronous_matches(true);
+
+  EXPECT_CALL(*client_.get(), IsOffTheRecord()).WillOnce(Return(false));
+  EXPECT_CALL(*client_.get(), SearchSuggestEnabled()).WillOnce(Return(true));
+
+  provider_->Start(input, false);
+  if (!provider_->done())
+    base::RunLoop().RunUntilIdle();
+
+  EXPECT_TRUE(provider_->matches().empty());
+  EXPECT_TRUE(provider_->done());
+}
+
+TEST_F(OnDeviceHeadProviderTest, RejectSynchronousRequest) {
+  AutocompleteInput input(base::UTF8ToUTF16("a"),
+                          metrics::OmniboxEventProto::OTHER,
+                          TestSchemeClassifier());
+  input.set_want_asynchronous_matches(false);
+
+  SetTestOnDeviceHeadServing();
+  provider_->Start(input, false);
+  if (!provider_->done())
+    base::RunLoop().RunUntilIdle();
+
+  EXPECT_TRUE(provider_->matches().empty());
+  EXPECT_TRUE(provider_->done());
+}
+
+TEST_F(OnDeviceHeadProviderTest, RejectIncognito) {
+  AutocompleteInput input(base::UTF8ToUTF16("a"),
+                          metrics::OmniboxEventProto::OTHER,
+                          TestSchemeClassifier());
+  input.set_want_asynchronous_matches(true);
+
+  EXPECT_CALL(*client_.get(), IsOffTheRecord()).WillOnce(Return(true));
+
+  SetTestOnDeviceHeadServing();
+  provider_->Start(input, false);
+  if (!provider_->done())
+    base::RunLoop().RunUntilIdle();
+
+  EXPECT_TRUE(provider_->matches().empty());
+  EXPECT_TRUE(provider_->done());
+}
+
+TEST_F(OnDeviceHeadProviderTest, NoMatches) {
+  AutocompleteInput input(base::UTF8ToUTF16("b"),
+                          metrics::OmniboxEventProto::OTHER,
+                          TestSchemeClassifier());
+  input.set_want_asynchronous_matches(true);
+
+  EXPECT_CALL(*client_.get(), IsOffTheRecord()).WillOnce(Return(false));
+  EXPECT_CALL(*client_.get(), SearchSuggestEnabled()).WillOnce(Return(true));
+
+  SetTestOnDeviceHeadServing();
+  provider_->Start(input, false);
+  if (!provider_->done())
+    base::RunLoop().RunUntilIdle();
+
+  EXPECT_TRUE(provider_->matches().empty());
+  EXPECT_TRUE(provider_->done());
+}
+
+TEST_F(OnDeviceHeadProviderTest, HasMatches) {
+  AutocompleteInput input(base::UTF8ToUTF16("M"),
+                          metrics::OmniboxEventProto::OTHER,
+                          TestSchemeClassifier());
+  input.set_want_asynchronous_matches(true);
+
+  EXPECT_CALL(*client_.get(), IsOffTheRecord()).WillOnce(Return(false));
+  EXPECT_CALL(*client_.get(), SearchSuggestEnabled()).WillOnce(Return(true));
+
+  SetTestOnDeviceHeadServing();
+  provider_->Start(input, false);
+  if (!provider_->done())
+    base::RunLoop().RunUntilIdle();
+
+  EXPECT_TRUE(provider_->done());
+  ASSERT_EQ(3U, provider_->matches().size());
+  EXPECT_EQ(base::UTF8ToUTF16("maps"), provider_->matches()[0].contents);
+  EXPECT_EQ(base::UTF8ToUTF16("mail"), provider_->matches()[1].contents);
+  EXPECT_EQ(base::UTF8ToUTF16("map"), provider_->matches()[2].contents);
+}
+
+TEST_F(OnDeviceHeadProviderTest, CancelInProgressRequest) {
+  AutocompleteInput input1(base::UTF8ToUTF16("g"),
+                           metrics::OmniboxEventProto::OTHER,
+                           TestSchemeClassifier());
+  input1.set_want_asynchronous_matches(true);
+  AutocompleteInput input2(base::UTF8ToUTF16("m"),
+                           metrics::OmniboxEventProto::OTHER,
+                           TestSchemeClassifier());
+  input2.set_want_asynchronous_matches(true);
+
+  EXPECT_CALL(*client_.get(), IsOffTheRecord()).WillRepeatedly(Return(false));
+  EXPECT_CALL(*client_.get(), SearchSuggestEnabled())
+      .WillRepeatedly(Return(true));
+
+  SetTestOnDeviceHeadServing();
+  provider_->Start(input1, false);
+  EXPECT_FALSE(provider_->done());
+  provider_->Start(input2, false);
+
+  if (!provider_->done())
+    base::RunLoop().RunUntilIdle();
+
+  EXPECT_TRUE(provider_->done());
+  ASSERT_EQ(3U, provider_->matches().size());
+  EXPECT_EQ(base::UTF8ToUTF16("maps"), provider_->matches()[0].contents);
+  EXPECT_EQ(base::UTF8ToUTF16("mail"), provider_->matches()[1].contents);
+  EXPECT_EQ(base::UTF8ToUTF16("map"), provider_->matches()[2].contents);
+}

components/omnibox/browser/on_device_head_serving.cc

-// Copyright (c) 2019 The Chromium Authors. All rights reserved.
+// Copyright 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.
 

components/omnibox/browser/on_device_head_serving.h

-// Copyright (c) 2019 The Chromium Authors. All rights reserved.
+// Copyright 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.
 

components/omnibox/browser/on_device_head_serving_unittest.cc

-// Copyright (c) 2019 The Chromium Authors. All rights reserved.
+// Copyright 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.