text:
#include "tensorflow/contrib/tensorrt/kernels/trt_engine_op.h"

#include <algorithm>
#include "tensorflow/contrib/tensorrt/convert/convert_nodes.h"
#include "tensorflow/contrib/tensorrt/convert/utils.h"
#include "tensorflow/contrib/tensorrt/log/trt_logger.h"
#include "tensorflow/contrib/tensorrt/resources/trt_resource_manager.h"
#include "tensorflow/contrib/tensorrt/resources/trt_resources.h"
#include "tensorflow/core/framework/graph_to_functiondef.h"
#include "tensorflow/core/lib/core/refcount.h"
#include "tensorflow/core/lib/strings/str_util.h"
#include "tensorflow/core/lib/strings/strcat.h"
#include "tensorflow/core/platform/logging.h"
#include "tensorflow/core/platform/stream_executor.h"
#include "tensorflow/core/platform/types.h"

#if GOOGLE_CUDA
#if GOOGLE_TENSORRT
#include "cuda/include/cuda_runtime_api.h"

namespace tensorflow {
namespace tensorrt {
static Logger logger;
using ::nvinfer1::IRuntime;
using ::tensorflow::strings::StrAppend;
using ::tensorflow::strings::StrCat;

// A helper class to call done() when destructed for asynchronous execution.
// Helps simultaneous execution of native and TRT engines.
class AsyncHelper : public tensorflow::core::RefCounted {
 public:
  AsyncHelper(tensorflow::AsyncOpKernel::DoneCallback done) { done_ = done; }
  ~AsyncHelper() override { done_(); }

 private:
  tensorflow::AsyncOpKernel::DoneCallback done_;
};

#define TYPECASE(dt, X, Y)                                                \
  case dt: {                                                              \
    return (void*)X->flat<tensorflow::EnumToDataType<dt>::Type>().data(); \
  }

void* GetTensorAddress(const Tensor* tensor_ptr) {
  auto tensor_type = tensor_ptr->dtype();
  switch (tensor_type) {
    TYPECASE(tensorflow::DT_FLOAT, tensor_ptr, dest_ptr);
    TYPECASE(tensorflow::DT_HALF, tensor_ptr, dest_ptr);
    TYPECASE(tensorflow::DT_INT8, tensor_ptr, dest_ptr);
    default: {
      LOG(ERROR) << "Unsupported Data type "
                 << tensorflow::DataTypeString(tensor_type);
      return nullptr;
    }
  }
}

tensorflow::Status TRTEngineOp::ConstructFunctionHandle(OpKernelContext* ctx) {
  VLOG(1) << "Constructing function handle";
  auto lib = ctx->function_library();
  if (lib == nullptr) {
    return tensorflow::errors::Internal("Context function library is null");
  }
  auto fdef = lib->GetFunctionLibraryDefinition()->Find(funcdef_name_);
  if (fdef == nullptr) {
    return tensorflow::errors::Internal("Native FunctionDef ", funcdef_name_,
                                        " can't be found in function library");
  }
  tensorflow::FunctionLibraryRuntime::InstantiateOptions inst_ops;
  inst_ops.overlay_lib = nullptr;
  inst_ops.state_handle = "";
  inst_ops.target = ctx->device()->name();
  native_func_ = 0;
  auto status = lib->Instantiate(funcdef_name_, AttrSlice(&fdef->attr()),
                                 inst_ops, &native_func_);
  if (!status.ok()) {
    LOG(ERROR) << " Instantiating native function " << funcdef_name_
               << " failed!";
  }
  return status;
}

TRTEngineOp::TRTEngineOp(OpKernelConstruction* context)
    : AsyncOpKernel(context) {
  // read serialized_engine
  OP_REQUIRES_OK(context,
                 context->GetAttr("serialized_segment", &serialized_segment_));
  OP_REQUIRES_OK(context,
                 context->GetAttr("workspace_size_bytes", &workspace_size_));
  OP_REQUIRES_OK(context, context->GetAttr("static_engine", &static_engine_));
  if (!static_engine_) {
    if (!segment_graph_.ParseFromString(serialized_segment_)) {
      LOG(ERROR) << "Parsing segment graph failed!";
      context->SetStatus(tensorflow::errors::InvalidArgument(
          "Failed to parse segment graphdef!"));
      return;
    }
    serialized_segment_.resize(0);
  }
  VLOG(1) << "Constructing " << name();
  string precision_string;
  OP_REQUIRES_OK(context,
                 context->GetAttr("precision_mode", &precision_string));
  string calibration_data;
  OP_REQUIRES_OK(context,
                 context->GetAttr("calibration_data", &calibration_data));
  OP_REQUIRES_OK(context,
                 context->GetAttr("segment_funcdef_name", &funcdef_name_));
  if (precision_string == "FP32") {
    precision_mode_ = convert::FP32MODE;
  } else if (precision_string == "FP16") {
    precision_mode_ = convert::FP16MODE;
  } else if (precision_string == "INT8") {
    precision_mode_ = convert::INT8MODE;
  }
  calibration_mode_ =
      (precision_mode_ == convert::INT8MODE && calibration_data.size() == 0);
  if (calibration_data.size()) {
    calibrator_.reset(new TRTInt8Calibrator(calibration_data));
    calibration_data.resize(0);
  }
  native_func_ = tensorflow::kInvalidHandle;
  OP_REQUIRES_OK(context, context->GetAttr("max_cached_engines_count",
                                           &max_cached_engines_));
  OP_REQUIRES_OK(context,
                 context->GetAttr("fixed_input_size", &fixed_input_size_));
  OP_REQUIRES_OK(context, context->GetAttr("cached_engine_batches",
                                           &cached_engine_batches_));
  std::sort(cached_engine_batches_.begin(), cached_engine_batches_.end());
  if (VLOG_IS_ON(1)) {
    string s("Engine Batches= ");
    for (auto i : cached_engine_batches_) {
      StrAppend(&s, i, " ");
    }
    VLOG(1) << s;
  }
}

void TRTEngineOp::ExecuteNativeSegment(tensorflow::OpKernelContext* ctx,
                                       AsyncHelper* helper) {
  if (!calibration_mode_) {
    VLOG(1) << "Executing native engine";
  }
  std::vector<Tensor> inputs;
  std::vector<Tensor>* outputs = new std::vector<Tensor>();
  if (native_func_ == tensorflow::kInvalidHandle) {
    auto status = ConstructFunctionHandle(ctx);
    if (!status.ok()) {
      LOG(ERROR) << "Couldn't construct function handle " << funcdef_name_;
      ctx->SetStatus(status);
      return;
    }
  }
  auto lib = ctx->function_library();
  tensorflow::FunctionLibraryRuntime::Options opts;
  opts.step_id = ctx->step_id();
  opts.rendezvous = ctx->rendezvous();
  opts.cancellation_manager = ctx->cancellation_manager();
  opts.runner = ctx->runner();
  for (int i = 0; i < ctx->num_inputs(); i++) {
    inputs.push_back(ctx->input(i));
  }
  helper->Ref();  // Increment count for calculating native graph
  VLOG(1) << "Executing native segment " << name();
  lib->Run(opts, native_func_, inputs, outputs,
           [ctx, outputs, helper](const tensorflow::Status& s) {
             tensorflow::core::ScopedUnref sc(helper);
             VLOG(1) << "Native Segment completed";
             if (!s.ok()) {
               ctx->SetStatus(s);
               return;
             }
             for (size_t t = 0; t < outputs->size(); ++t) {
               ctx->set_output(t, outputs->at(t));
             }
             delete outputs;
           });
}

void TRTEngineOp::ExecuteCalibration(tensorflow::OpKernelContext* ctx,
                                     AsyncHelper* helper) {
  helper->Ref();
  tensorflow::core::ScopedUnref sc(helper);
  // TODO(aaroey): remove the ResourceMgr singleton.
  auto trt_rm = TRTResourceManager::instance();
  auto res_mgr = trt_rm->getManager("TRTCalibration");
  TRTCalibrationResource* calib_res = nullptr;
  auto status = res_mgr->LookupOrCreate(
      funcdef_name_, "Calibrator", &calib_res,
      {[ctx, this](TRTCalibrationResource** cr) -> tensorflow::Status {
        return this->AllocateCalibrationResources(ctx, cr);
      }});
  if (!status.ok()) {
    ctx->SetStatus(status);
    return;
  }
  int num_inputs = ctx->num_inputs();
  // Pass input data to calibrator
  std::unordered_map<string, void*> input_data;
  for (int i = 0; i < num_inputs; i++) {
    const Tensor& t = ctx->input(i);
    void* data_address = GetTensorAddress(&t);
    if (data_address == nullptr) {
      ctx->SetStatus(tensorflow::errors::InvalidArgument(
          "Unsupported data type encountered in input ", i));
      return;
    }
    // Check the allocated buffer is sufficient for input
    const auto device_tensor = dev_tensors_.at(i).AccessTensor(ctx);
    CHECK_EQ(t.TotalBytes(), device_tensor->TotalBytes());
    input_data.emplace(StrCat(kInputPHName, i), data_address);
  }
  VLOG(2) << "Filled map for sending";
  // copied from cuda_kernel_helper since it seems only valid in *.cu.cc files
  const cudaStream_t* stream = CHECK_NOTNULL(
      reinterpret_cast<const cudaStream_t*>(ctx->op_device_context()
                                                ->stream()
                                                ->implementation()
                                                ->CudaStreamMemberHack()));
  calib_res->calibrator_->setBatch(input_data, *stream);
  VLOG(2) << "Passed calibration data";
  ExecuteNativeSegment(ctx, helper);
}

int TRTEngineOp::GetEngineBatch(tensorflow::OpKernelContext* ctx) {
  int num_batch = ctx->input(0).shape().dim_size(0);
  int smallest_engine = 0;
  for (const auto i : cached_engine_batches_) {
    if (i >= num_batch) {
      smallest_engine = i;
      break;
    }
  }
  // TODO(sami): Need an LRU here
  if (smallest_engine == 0) {
    if (max_cached_engines_ > cached_engine_batches_.size()) {
      smallest_engine = num_batch;
      cached_engine_batches_.push_back(num_batch);
      VLOG(1) << "Running with batch size " << num_batch;
    } else {
      string s("Engine buffer is full. buffer limit= ");
      StrAppend(&s, max_cached_engines_, ", current entries= ");
      for (auto i : cached_engine_batches_) StrAppend(&s, i, ", ");
      StrAppend(&s, "Requested batch= ", num_batch);
      LOG(ERROR) << s;
      ctx->SetStatus(tensorflow::errors::ResourceExhausted(
          "Requested batch size is not available and engine cache is full"));
      return -1;
    }
  }
  return smallest_engine;
}

void TRTEngineOp::ComputeAsync(tensorflow::OpKernelContext* ctx,
                               tensorflow::AsyncOpKernel::DoneCallback done) {
  auto helper = new AsyncHelper(done);
  tensorflow::core::ScopedUnref sc(helper);
  if (calibration_mode_) {
    ExecuteCalibration(ctx, helper);
    return;
  }
  const int smallest_engine = GetEngineBatch(ctx);
  if (smallest_engine < 0) return;  // GetEngineBatch already set the status.

  const int num_batch = ctx->input(0).shape().dim_size(0);
  auto& engine_ctx_pair = GetEngine(smallest_engine, ctx);
  auto& trt_engine_ptr = engine_ctx_pair.first;
  if (!trt_engine_ptr) {
    LOG(WARNING) << "Engine retrieval for batch size " << num_batch
                 << " failed Running native segment";
    ExecuteNativeSegment(ctx, helper);
    return;
  }

  const int num_binding = ctx->num_inputs() + ctx->num_outputs();
  std::vector<void*> buffers(num_binding);
  for (int i = 0; i < ctx->num_inputs(); i++) {
    const string inp_name = StrCat(kInputPHName, i);
    const size_t binding_index =
        trt_engine_ptr->getBindingIndex(inp_name.c_str());

    const Tensor& input_tensor = ctx->input(i);
    const TensorShape& input_shape = input_tensor.shape();
    if (num_batch != input_shape.dim_size(0)) {
      LOG(ERROR) << "input data inconsistent batch size";
      ctx->SetStatus(tensorflow::errors::FailedPrecondition(
          "Different batch sizes between input tensors"));
      return;
    }
    auto dtype = trt_engine_ptr->getBindingDataType(binding_index);
    switch (dtype) {
      case nvinfer1::DataType::kFLOAT:
        buffers[binding_index] = (void*)(input_tensor.flat<float>().data());
        break;
      case nvinfer1::DataType::kHALF:
        LOG(ERROR) << "FP16 inputs are not supported yet!";
        ctx->SetStatus(tensorflow::errors::InvalidArgument(
            "FP16 inputs are not supported!"));
        return;
      case nvinfer1::DataType::kINT8:
        LOG(ERROR) << "INT8 inputs are not supported yet!";
        ctx->SetStatus(tensorflow::errors::InvalidArgument(
            "INT8 inputs are not supported!"));
        return;
      default:
        LOG(ERROR) << "Unknown TRT data type: " << int(dtype);
        ctx->SetStatus(tensorflow::errors::InvalidArgument(
            "Unknown ouput TRT data type! ", static_cast<int>(dtype)));
        return;
    }
  }

  for (int i = 0; i < ctx->num_outputs(); i++) {
    // Create an output tensor
    const string output_name = StrCat(kOutputPHName, i);
    const size_t binding_index = trt_engine_ptr->getBindingIndex(
        output_name.c_str());
    Tensor* output_tensor = nullptr;

    TensorShape output_shape;
    if (binding_index != -1) {
      auto dims = trt_engine_ptr->getBindingDimensions(binding_index);
      std::vector<int> trt_shape(dims.nbDims + 1);
      trt_shape[0] = num_batch;
      for (int j = 0; j < dims.nbDims; j++) trt_shape[j + 1] = dims.d[j];
      OP_REQUIRES_OK(
          ctx, TensorShapeUtils::MakeShape(trt_shape.data(), trt_shape.size(),
                                           &output_shape));
    } else {
      LOG(ERROR) << "output node not found, at " << output_name;
      ctx->SetStatus(tensorflow::errors::Internal("output ", output_name,
                                                  " couldn't be found!"));
      return;
    }
    auto status = ctx->allocate_output(i, output_shape, &output_tensor);
    if (!status.ok()) {
      LOG(ERROR) << "Allocating output failed with " << status;
      ctx->SetStatus(status);
      return;
    }
    auto dtype = trt_engine_ptr->getBindingDataType(binding_index);
    switch (dtype) {
      case nvinfer1::DataType::kFLOAT:
        buffers[binding_index] =
            reinterpret_cast<void*>(output_tensor->flat<float>().data());
        break;
      case nvinfer1::DataType::kHALF:
        LOG(ERROR) << "half size is not supported yet!";
        ctx->SetStatus(tensorflow::errors::InvalidArgument(
            "Half outputs are not supported!"));
        return;
      case nvinfer1::DataType::kINT8:
        LOG(ERROR) << "int8 is not supported yet!";
        ctx->SetStatus(tensorflow::errors::InvalidArgument(
            "INT8 outputs are not supported!"));
        return;
      default:
        LOG(ERROR) << "Unknown TRT data type: " << static_cast<int>(dtype);
        ctx->SetStatus(tensorflow::errors::InvalidArgument(
            "Unsupported output data type! ", int(dtype)));
        return;
    }
  }
  // copied from cuda_kernel_helper since it seems only valid in *.cu.cc files
  const cudaStream_t* stream = CHECK_NOTNULL(
      reinterpret_cast<const cudaStream_t*>(ctx->op_device_context()
                                                ->stream()
                                                ->implementation()
                                                ->CudaStreamMemberHack()));

  // TODO(jie): trt enqueue does not return error
  auto& trt_execution_context_ptr = engine_ctx_pair.second;
  auto ret = trt_execution_context_ptr->enqueue(num_batch, &buffers[0], *stream,
                                                nullptr);
  if (!ret) {
    LOG(ERROR) << "Failed to enqueue batch for TRT engine: " << name();
    ctx->SetStatus(tensorflow::errors::Internal(
        "Failed to enqueue batch for TRT engine: ", name()));
  }
  // sync should be done by TF.
}

TRTEngineOp::~TRTEngineOp() {
  // We need to manually destroy the engine and execution context before
  // the allocator is destructed.
  for (auto& eng : engine_map_) {
    eng.second.first.reset();
    eng.second.second.reset();
  }
  allocator_.reset();
}

nvinfer1::IGpuAllocator* TRTEngineOp::GetAllocator(OpKernelContext* ctx) {
  if (allocator_) return allocator_.get();
  auto device = ctx->device();
  auto alloc = device->GetAllocator(tensorflow::AllocatorAttributes());
  if (!alloc) {
    LOG(ERROR) << "Can't find device allocator for gpu device "
               << device->name();
    ctx->SetStatus(tensorflow::errors::Internal(
        "Can't get device allocator for device ", device->name()));
    return nullptr;
  }
  allocator_.reset(new TRTDeviceAllocator(alloc));
  return allocator_.get();
}

TRTEngineOp::EngineCtxPair& TRTEngineOp::GetEngine(int batch_size,
                                                  OpKernelContext* ctx) {
  static EngineCtxPair null_pair = {
    TrtUniquePtrType<nvinfer1::ICudaEngine>(nullptr),
    TrtUniquePtrType<nvinfer1::IExecutionContext>(nullptr)};
  // TODO(sami): This method needs to be re-written to use resource manager and
  // with LRU mechanism option.
  tensorflow::mutex_lock lock(engine_mutex_);

  if (static_engine_) {
    if (engine_map_.size()) {
      if (engine_map_.begin()->first >= batch_size) {
        return engine_map_.begin()->second;
      }
      return null_pair;
    }
    TrtUniquePtrType<IRuntime> infer(nvinfer1::createInferRuntime(logger));
#if NV_TENSORRT_MAJOR > 3
    auto allocator = GetAllocator(ctx);
    if (allocator == nullptr) {
      // GetAllocator already set the Status.
      return null_pair;
    }
    infer->setGpuAllocator(allocator);
#endif
    TrtUniquePtrType<nvinfer1::ICudaEngine> static_engine(
        infer->deserializeCudaEngine(serialized_segment_.c_str(),
                                     serialized_segment_.size(), nullptr));
    auto raw_static_engine = static_engine.get();
    const auto max_batch_size = raw_static_engine->getMaxBatchSize();
    engine_map_[max_batch_size] = {
      std::move(static_engine),
      TrtUniquePtrType<nvinfer1::IExecutionContext>(
          raw_static_engine->createExecutionContext())};
    // Runtime is safe to delete after engine creation
    serialized_segment_.clear();
    if (max_batch_size < batch_size) return null_pair;
    return engine_map_.at(max_batch_size);
  }  // static_engine_

  // Handle the dynamic engine case.
  auto engine_it = engine_map_.find(batch_size);
  if (engine_it == engine_map_.end() &&
      engine_map_.size() < (size_t)max_cached_engines_) {
    nvinfer1::IGpuAllocator* allocator = nullptr;
#if NV_TENSORRT_MAJOR > 3
    allocator = GetAllocator(ctx);
    if (allocator == nullptr) {
      // GetAllocator already set the Status.
      return null_pair;
    }
#endif
    std::vector<tensorflow::PartialTensorShape> shapes;
    for (int i = 0; i < ctx->num_inputs(); ++i) {
      shapes.emplace_back(ctx->input(i).shape());
    }
    TrtUniquePtrType<nvinfer1::ICudaEngine> engine;
    bool convert_successfully = false;
    VLOG(0) << name() << " Constructing a new engine with batch size "
            << batch_size;
    // Up to this point, calibrator_ can never be empty, since otherwise it
    // means calibration_mode_ is true and this path won't get executed.
    auto status = convert::ConvertGraphDefToEngine(
        segment_graph_, precision_mode_, batch_size, workspace_size_, shapes,
        &logger, allocator, calibrator_.get(), &engine, &convert_successfully);
    if (!status.ok()) {
      if (convert_successfully) {
        // This means it fail to build the engine even when the network is built
        // successfully, probably due to internal issues. In this case we don't
        // retry in the future.
        engine_map_[batch_size] = {nullptr, nullptr};
      }
      LOG(ERROR) << "Engine creation for batch size " << batch_size
                 << " failed " << status;
      ctx->SetStatus(tensorflow::errors::Internal("Engine creation failed!"));
      return null_pair;
    }
    VLOG(1) << "Conversion is done";
    TrtUniquePtrType<nvinfer1::IExecutionContext> exec_context(
        engine->createExecutionContext());
    engine_map_[batch_size] = {std::move(engine), std::move(exec_context)};
  }
  return engine_map_.at(batch_size);
}

tensorflow::Status TRTEngineOp::AllocateCalibrationResources(
    tensorflow::OpKernelContext* ctx, TRTCalibrationResource** cr) {
  auto cres = new TRTCalibrationResource();
  *cr = cres;
  // Get the allocator.
  auto alloc = ctx->device()->GetAllocator(tensorflow::AllocatorAttributes());
  if (!alloc) {
    LOG(WARNING) << "Can't get device allocator will not be able to "
                    "allocate memory from TensorFlow memory pool";
    cres->allocator_.reset(new TRTCudaAllocator);
  } else {
    cres->allocator_.reset(new TRTDeviceAllocator(alloc));
  }
  // Get the input shapes.
  const int batch_size = ctx->input(0).dim_size(0);
  const int num_inputs = ctx->num_inputs();
  std::vector<tensorflow::PartialTensorShape> shapes;
  dev_tensors_.resize(num_inputs);
  VLOG(1) << " Constructing calibrator";
  for (int i = 0; i < num_inputs; i++) {
    // allocate workspace on device for inputs
    const tensorflow::Tensor& t = ctx->input(i);
    shapes.emplace_back(t.shape());
    Tensor* device_tensor;
    TF_RETURN_IF_ERROR(ctx->allocate_persistent(
        t.dtype(), t.shape(), &dev_tensors_.at(i), &device_tensor));
    CHECK_EQ(t.TotalBytes(), device_tensor->TotalBytes());
    void* device_address = GetTensorAddress(device_tensor);
    if (device_address == nullptr) {
      return tensorflow::errors::InvalidArgument(
          "Unsupported data type encountered in input ", i);
    }
    device_buffers_.emplace(
        StrCat(kInputPHName, i),
        std::pair<void*, size_t>(device_address, device_tensor->TotalBytes()));
  }
  cres->calibrator_.reset(
      new TRTInt8Calibrator(device_buffers_, batch_size, name()));
  const string label(name());
  auto segment_graph = &segment_graph_;
  const int cuda_gpu_id = ctx->device()->tensorflow_gpu_device_info()->gpu_id;
  if (cuda_gpu_id < 0) {
    LOG(ERROR) << "Can't get gpu_device_info from context->device()";
    return tensorflow::errors::InvalidArgument(
        "Context->device doesn't contain device info!");
  }
  const int64 workspace_size_bytes = workspace_size_;
  cres->thr_.reset(new std::thread([cres, label, segment_graph, shapes,
                                    cuda_gpu_id, workspace_size_bytes]() {
    VLOG(0) << "Starting calibration thread on device " << cuda_gpu_id
            << ", Calibration Resource @ " << cres;
    auto err = cudaSetDevice(cuda_gpu_id);
    if (err != cudaSuccess) {
      // TODO(aaroey): should return error here.
      LOG(ERROR) << "Couldn't set cuda device to " << cuda_gpu_id
                 << " in calibration thread";
    }
    // ConvertGraphDefToEngine() will try to build the engine. This thread
    // will loop inside buildCudaEngine() consuming the calibration data
    // that is set by the TF op, and drive the builder until calibrator returns
    // false. Engine is discarded after calibration table is generated
    //
    // TODO(aaroey): maybe setting the max batch size using the python
    // calibration wrapper class.
    auto s = convert::ConvertGraphDefToEngine(
        *segment_graph, convert::INT8MODE, cres->calibrator_->getBatchSize(),
        workspace_size_bytes, shapes, &cres->logger_, cres->allocator_.get(),
        cres->calibrator_.get(), &cres->engine_,
        /*convert_successfully=*/nullptr);
    if (!s.ok()) {
      LOG(ERROR) << "Calibration failed: " << s;
      cres->calibrator_->setDone();  // Ignore further pushes
    }
    VLOG(1) << "Calibration loop terminated " << label;
  }));
  VLOG(1) << "initialized calibrator resource";
  return tensorflow::Status::OK();
}

REGISTER_KERNEL_BUILDER(Name("TRTEngineOp").Device(DEVICE_GPU), TRTEngineOp);

}  // namespace tensorrt
}  // namespace tensorflow

#endif  // GOOGLE_TENSORRT
#endif  // GOOGLE_CUDA


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

import 'chrome://diagnostics/input_list.js';
import 'chrome://diagnostics/strings.m.js';
import 'chrome://resources/mojo/mojo/public/js/mojo_bindings_lite.js';

import {DiagnosticsBrowserProxyImpl} from 'chrome://diagnostics/diagnostics_browser_proxy.js';
import {NavigationView} from 'chrome://diagnostics/diagnostics_types.js';
import {fakeKeyboards, fakeTouchDevices} from 'chrome://diagnostics/fake_data.js';
import {FakeInputDataProvider} from 'chrome://diagnostics/fake_input_data_provider.js';
import {InputCardElement} from 'chrome://diagnostics/input_card.js';
import {ConnectionType, KeyboardInfo, MechanicalLayout, NumberPadPresence, PhysicalLayout, TopRightKey, TopRowKey, TouchDeviceInfo, TouchDeviceType} from 'chrome://diagnostics/input_data_provider.mojom-webui.js';
import {InputListElement} from 'chrome://diagnostics/input_list.js';
import {setInputDataProviderForTesting} from 'chrome://diagnostics/mojo_interface_provider.js';
import {loadTimeData} from 'chrome://resources/js/load_time_data.m.js';
import {PromiseResolver} from 'chrome://resources/js/promise_resolver.js';
import {flushTasks} from 'chrome://webui-test/polymer_test_util.js';
import {eventToPromise} from 'chrome://webui-test/test_util.js';

import {assertArrayEquals, assertEquals, assertFalse, assertTrue} from 'chrome://webui-test/chromeos/chai_assert.js';
import {isVisible} from '../test_util.js';

import {TestDiagnosticsBrowserProxy} from './test_diagnostics_browser_proxy.js';

suite('inputListTestSuite', function() {
  /** @type {?InputListElement} */
  let inputListElement = null;

  /** @type {?FakeInputDataProvider} */
  let provider = null;

  /** @type {?TestDiagnosticsBrowserProxy} */
  let diagnosticsBrowserProxy = null;

  suiteSetup(() => {
    provider = new FakeInputDataProvider();
    setInputDataProviderForTesting(provider);

    diagnosticsBrowserProxy = new TestDiagnosticsBrowserProxy();
    DiagnosticsBrowserProxyImpl.setInstance(diagnosticsBrowserProxy);
  });

  setup(() => {
    document.body.innerHTML = '';
  });

  teardown(() => {
    inputListElement.remove();
    inputListElement = null;
    provider.reset();
  });

  function initializeInputList(
      keyboards = fakeKeyboards, touchDevices = fakeTouchDevices) {
    assertFalse(!!inputListElement);
    provider.setFakeConnectedDevices(keyboards, touchDevices);

    // Add the input list to the DOM.
    inputListElement =
        /** @type {!InputListElement} */ (document.createElement('input-list'));
    assertTrue(!!inputListElement);
    document.body.appendChild(inputListElement);

    return flushTasks();
  }

  /** @return {!InputCardElement} */
  function getCardByDeviceType(deviceType) {
    const card = inputListElement.shadowRoot.querySelector(
        `input-card[device-type="${deviceType}"]`);
    return /** @type {!InputCardElement} */ (card);
  }

  /**
   * openTouchscreenTester is a helper function for some boilerplate code. It
   * clicks the test button of a touchscreen, then clicks the start testing
   * button and makes sure canvas dialog is open.
   * @param touchscreensToInitialize The touchscreens to initialize.
   * @param isClamshellMode If the tester is opened in clamshell mode.
   * @returns The touchscreenTester.
   */
  async function openTouchscreenTester(
      touchscreensToInitialize, isClamshellMode = true) {
    await initializeInputList([], touchscreensToInitialize);
    if (isClamshellMode) {
      provider.setStartTesterWithClamshellMode();
    } else {
      provider.setStartTesterWithTabletMode();
    }

    const resolver = new PromiseResolver();
    const touchscreenTester =
        inputListElement.shadowRoot.querySelector('touchscreen-tester');
    const introDialog = touchscreenTester.getDialog('intro-dialog');

    // Mock requestFullscreen function since this API can only be initiated by a
    // user gesture.
    introDialog.requestFullscreen = () => {
      resolver.resolve();
    };

    const testButton = getCardByDeviceType('touchscreen')
                           .shadowRoot.querySelector('cr-button');
    assertTrue(!!testButton);
    testButton.click();
    await flushTasks();
    assertTrue(introDialog.open);

    const getStartedButton = introDialog.querySelector('cr-button');
    getStartedButton.click();
    await flushTasks();
    const canvasDialog = touchscreenTester.getDialog('canvas-dialog');
    assertTrue(canvasDialog.open);

    return touchscreenTester;
  }

  test('InputListPopulated', async () => {
    await initializeInputList();
    assertEquals(
        fakeKeyboards[0].name, getCardByDeviceType('keyboard').devices[0].name);
    assertEquals(
        fakeTouchDevices[0].name,
        getCardByDeviceType('touchpad').devices[0].name);
    assertEquals(
        fakeTouchDevices[1].name,
        getCardByDeviceType('touchscreen').devices[0].name);
  });

  test('KeyboardAddAndRemove', async () => {
    /** @type {!KeyboardInfo} */
    const fakeKeyboard = {
      id: 4,
      connectionType: ConnectionType.kUsb,
      name: 'Sample USB keyboard',
      physicalLayout: PhysicalLayout.kUnknown,
      mechanicalLayout: MechanicalLayout.kUnknown,
      hasAssistantKey: false,
      numberPadPresent: NumberPadPresence.kUnknown,
      topRowKeys: [
        TopRowKey.kBack,
        TopRowKey.kForward,
        TopRowKey.kRefresh,
        TopRowKey.kFullscreen,
        TopRowKey.kOverview,
        TopRowKey.kScreenBrightnessDown,
        TopRowKey.kScreenBrightnessUp,
        TopRowKey.kVolumeMute,
        TopRowKey.kVolumeDown,
        TopRowKey.kVolumeUp,
      ],
      topRightKey: TopRightKey.kUnknown,
    };
    await initializeInputList();
    const keyboardCard = getCardByDeviceType('keyboard');

    provider.addFakeConnectedKeyboard(fakeKeyboard);
    await flushTasks();
    assertEquals(2, keyboardCard.devices.length);
    assertEquals(fakeKeyboards[0].name, keyboardCard.devices[0].name);
    assertEquals(fakeKeyboard.name, keyboardCard.devices[1].name);

    provider.removeFakeConnectedKeyboardById(fakeKeyboard.id);
    await flushTasks();
    assertEquals(1, keyboardCard.devices.length);
    assertEquals(fakeKeyboards[0].name, keyboardCard.devices[0].name);
  });

  test('KeyboardTesterShow', async () => {
    await initializeInputList([fakeKeyboards[0]], []);
    const testButton =
        getCardByDeviceType('keyboard').shadowRoot.querySelector('cr-button');
    assertTrue(!!testButton);
    testButton.click();
    await flushTasks();

    const keyboardTester =
        inputListElement.shadowRoot.querySelector('keyboard-tester');
    assertTrue(keyboardTester.isOpen());
  });

  test('ShowToastIfKeyboardDisconnectedDuringTest', async () => {
    await initializeInputList([fakeKeyboards[0]], []);
    const testButton =
        getCardByDeviceType('keyboard').shadowRoot.querySelector('cr-button');
    assertTrue(!!testButton);
    testButton.click();
    await flushTasks();

    const keyboardTester =
        inputListElement.shadowRoot.querySelector('keyboard-tester');
    assertTrue(keyboardTester.isOpen());
    const showToastEvent = eventToPromise('show-toast', inputListElement);
    // Remove keyboard while tester is open.
    provider.removeFakeConnectedKeyboardById(fakeKeyboards[0].id);
    await flushTasks();
    // Verify that the custom event was dispatched
    const e = await showToastEvent;
    assertEquals(
        e.detail.message, loadTimeData.getString('deviceDisconnected'));

    // Verify that tester is closed.
    assertFalse(keyboardTester.isOpen());

    // Verify that key events are no longer blocked by keyboard-tester.
    let keyEventReceived = false;
    inputListElement.addEventListener('keydown', () => keyEventReceived = true);
    const keyDownEvent = eventToPromise('keydown', inputListElement);
    keyboardTester.dispatchEvent(new KeyboardEvent(
        'keydown', {bubbles: true, key: 'A', composed: true}));
    await keyDownEvent;
    assertTrue(keyEventReceived);
  });

  test('TouchpadAddAndRemove', async () => {
    /** @type {!TouchDeviceInfo} */
    const fakeTouchpad = {
      id: 4,
      connectionType: ConnectionType.kUsb,
      type: TouchDeviceType.kPointer,
      name: 'Sample USB touchpad',
    };
    await initializeInputList();
    const touchpadCard = getCardByDeviceType('touchpad');

    provider.addFakeConnectedTouchDevice(fakeTouchpad);
    await flushTasks();
    assertEquals(2, touchpadCard.devices.length);
    assertEquals(fakeTouchDevices[0].name, touchpadCard.devices[0].name);
    assertEquals(fakeTouchpad.name, touchpadCard.devices[1].name);

    provider.removeFakeConnectedTouchDeviceById(fakeTouchpad.id);
    await flushTasks();
    assertEquals(1, touchpadCard.devices.length);
    assertEquals(fakeTouchDevices[0].name, touchpadCard.devices[0].name);
  });

  test('TouchpadTesterShowAndClose', async () => {
    await initializeInputList([], [fakeTouchDevices[0]]);
    const touchpadTester =
        inputListElement.shadowRoot.querySelector('touchpad-tester');
    assertTrue(!!touchpadTester);
    assertFalse(touchpadTester.isOpen());

    const testButton =
        getCardByDeviceType('touchpad').shadowRoot.querySelector('cr-button');
    assertTrue(!!testButton);
    testButton.click();
    await flushTasks();

    assertTrue(touchpadTester.isOpen());
    assertDeepEquals(fakeTouchDevices[0], touchpadTester.touchpad);
  });

  test('TouchscreenAddAndRemove', async () => {
    /** @type {!TouchDeviceInfo} */
    const fakeTouchscreen = {
      id: 4,
      connectionType: ConnectionType.kUsb,
      type: TouchDeviceType.kDirect,
      name: 'Sample USB touchscreen',
    };
    await initializeInputList();
    const touchscreenCard = getCardByDeviceType('touchscreen');

    provider.addFakeConnectedTouchDevice(fakeTouchscreen);
    await flushTasks();
    assertEquals(3, touchscreenCard.devices.length);
    assertEquals(fakeTouchDevices[1].name, touchscreenCard.devices[0].name);
    assertEquals(fakeTouchscreen.name, touchscreenCard.devices[2].name);

    provider.removeFakeConnectedTouchDeviceById(fakeTouchscreen.id);
    await flushTasks();
    assertEquals(2, touchscreenCard.devices.length);
    assertEquals(fakeTouchDevices[1].name, touchscreenCard.devices[0].name);
  });

  test('TouchscreenTesterShowAndClose', async () => {
    await initializeInputList([], [fakeTouchDevices[1]]);
    provider.setStartTesterWithClamshellMode();

    const resolver = new PromiseResolver();
    let requestFullscreenCalled = 0;

    const touchscreenTester =
        inputListElement.shadowRoot.querySelector('touchscreen-tester');
    const introDialog = touchscreenTester.getDialog('intro-dialog');

    // Mock requestFullscreen function since this API can only be initiated by a
    // user gesture.
    introDialog.requestFullscreen = () => {
      requestFullscreenCalled++;
      resolver.resolve();
    };

    let eventTrackerRemoveAllCalled = 0;
    touchscreenTester.getEventTracker().removeAll = () => {
      eventTrackerRemoveAllCalled++;
      resolver.resolve();
    };

    // A11y touch passthrough state is false by default.
    assertFalse(provider.getA11yTouchPassthroughState());

    const testButton = getCardByDeviceType('touchscreen')
                           .shadowRoot.querySelector('cr-button');
    assertTrue(!!testButton);
    testButton.click();
    await flushTasks();
    assertEquals(1, requestFullscreenCalled);
    assertTrue(introDialog.open);
    assertEquals(
        /*expectedMoveAppToTestingScreenCalled=*/ 1,
        provider.getMoveAppToTestingScreenCalled());

    // Click get stated button to open canvas dialog.
    const getStartedButton = introDialog.querySelector('cr-button');
    getStartedButton.click();
    await flushTasks();
    assertFalse(introDialog.open);
    assertTrue(provider.getA11yTouchPassthroughState());

    const canvasDialog = touchscreenTester.getDialog('canvas-dialog');
    assertTrue(canvasDialog.open);

    const fullscreenChangeEvent = eventToPromise('fullscreenchange', document);
    document.dispatchEvent(new Event('fullscreenchange'));
    await fullscreenChangeEvent;

    assertFalse(canvasDialog.open);
    assertEquals(
        /*expectedMoveAppBackToPreviousScreenCalled=*/ 1,
        provider.getMoveAppBackToPreviousScreenCalled());
    assertEquals(
        /*expectedEventTrackerRemoveAllCalled=*/ 1,
        eventTrackerRemoveAllCalled);
    assertFalse(provider.getA11yTouchPassthroughState());
  });

  test('TouchscreenTesterShowAndCloseInTabletMode', async () => {
    await openTouchscreenTester(
        /*touchscreensToInitialize=*/[fakeTouchDevices[1]]);

    const resolver = new PromiseResolver();
    let exitFullscreenCalled = 0;
    // Mock exitFullscreen call.
    document.exitFullscreen = () => {
      exitFullscreenCalled++;
      resolver.resolve();
    };

    provider.startTabletMode();
    await flushTasks();

    const keyEvent = eventToPromise('keydown', window);
    window.dispatchEvent(new KeyboardEvent('keydown', {key: 'AudioVolumeUp'}));
    await keyEvent;
    assertEquals(1, exitFullscreenCalled);
  });

  test('StartTouchscreenTesterWithClamshellMode', async () => {
    const touchscreenTester = await openTouchscreenTester(
        /*touchscreensToInitialize=*/[fakeTouchDevices[1]]);

    assertFalse(touchscreenTester.getIsTabletMode());

    provider.startTabletMode();
    await flushTasks();

    assertTrue(touchscreenTester.getIsTabletMode());
    assertTrue(touchscreenTester.getDialog('canvas-dialog').open);
  });

  test('StartTouchscreenTesterWithTabletMode', async () => {
    const touchscreenTester = await openTouchscreenTester(
        /*touchscreensToInitialize=*/[fakeTouchDevices[1]],
        /*isClamshellMode=*/ false);

    assertTrue(touchscreenTester.getIsTabletMode());

    provider.endTabletMode();
    await flushTasks();

    assertFalse(touchscreenTester.getIsTabletMode());
    assertTrue(touchscreenTester.getDialog('canvas-dialog').open);
  });

  test('OnInternalDisplayPowerStateChanged', async () => {
    await initializeInputList([], [fakeTouchDevices[1]]);

    assertTrue(getCardByDeviceType('touchscreen').devices[0].testable);

    provider.setInternalDisplayPowerOff();
    await flushTasks();

    assertFalse(getCardByDeviceType('touchscreen').devices[0].testable);

    provider.setInternalDisplayPowerOn();
    await flushTasks();

    assertTrue(getCardByDeviceType('touchscreen').devices[0].testable);
  });

  test('InternalDisplayPowerOffWhileTouchscreenTesterIsRunning', async () => {
    const touchscreenTester = await openTouchscreenTester(
        /*touchscreensToInitialize=*/[fakeTouchDevices[1]]);

    // Internal display power off.
    provider.setInternalDisplayPowerOff();
    await flushTasks();

    // Tester is expecetd to exit.
    assertFalse(touchscreenTester.getDialog('canvas-dialog').open);
  });

  test('TouchscreenDisconnectedWhileTesterIsRunning', async () => {
    const touchscreenTester = await openTouchscreenTester(
        /*touchscreensToInitialize=*/[fakeTouchDevices[1]]);

    // Remove touchscreen.
    provider.removeFakeConnectedTouchDeviceById(fakeTouchDevices[1].id);
    await flushTasks();

    // Tester is expecetd to exit.
    assertFalse(touchscreenTester.getDialog('canvas-dialog').open);
  });

  test('EmptySectionsHidden', async () => {
    await initializeInputList([], []);
    assertFalse(isVisible(getCardByDeviceType('keyboard')));
    assertFalse(isVisible(getCardByDeviceType('touchpad')));
    assertFalse(isVisible(getCardByDeviceType('touchscreen')));

    provider.addFakeConnectedTouchDevice(fakeTouchDevices[1]);
    await flushTasks();
    assertTrue(!!getCardByDeviceType('touchscreen'));
    assertFalse(isVisible(getCardByDeviceType('keyboard')));
    assertFalse(isVisible(getCardByDeviceType('touchpad')));
    assertTrue(isVisible(getCardByDeviceType('touchscreen')));

    provider.removeFakeConnectedTouchDeviceById(fakeTouchDevices[1].id);
    provider.addFakeConnectedTouchDevice(fakeTouchDevices[0]);
    await flushTasks();
    assertFalse(isVisible(getCardByDeviceType('keyboard')));
    assertTrue(isVisible(getCardByDeviceType('touchpad')));
    assertFalse(isVisible(getCardByDeviceType('touchscreen')));

    provider.removeFakeConnectedTouchDeviceById(fakeTouchDevices[0].id);
    provider.addFakeConnectedKeyboard(fakeKeyboards[0]);
    await flushTasks();
    assertTrue(isVisible(getCardByDeviceType('keyboard')));
    assertFalse(isVisible(getCardByDeviceType('touchpad')));
    assertFalse(isVisible(getCardByDeviceType('touchscreen')));

    provider.removeFakeConnectedKeyboardById(fakeKeyboards[0].id);
    await flushTasks();
    assertFalse(isVisible(getCardByDeviceType('keyboard')));
    assertFalse(isVisible(getCardByDeviceType('touchpad')));
    assertFalse(isVisible(getCardByDeviceType('touchscreen')));
  });

  test('RecordNavigationCalled', async () => {
    await initializeInputList();
    inputListElement.onNavigationPageChanged({isActive: false});
    await flushTasks();

    assertEquals(0, diagnosticsBrowserProxy.getCallCount('recordNavigation'));

    diagnosticsBrowserProxy.setPreviousView(NavigationView.SYSTEM);
    inputListElement.onNavigationPageChanged({isActive: true});

    await flushTasks();
    assertEquals(1, diagnosticsBrowserProxy.getCallCount('recordNavigation'));
    assertArrayEquals(
        [NavigationView.SYSTEM, NavigationView.INPUT],
        /** @type {!Array<!NavigationView>} */
        (diagnosticsBrowserProxy.getArgs('recordNavigation')[0]));
  });
});

redpajama_set_name: RedPajamaGithub