// Copyright (c) 2018, the Dart project authors.  Please see the AUTHORS file
// for details. All rights reserved. Use of this source code is governed by a
// BSD-style license that can be found in the LICENSE file.

import 'dart:async';
import 'dart:typed_data';

import 'package:kernel/ast.dart' as ir;
import 'package:kernel/binary/ast_from_binary.dart' as ir;
import 'package:kernel/binary/ast_to_binary.dart' as ir;

import '../../compiler_api.dart' as api;
import '../common/codegen.dart';
import '../common/tasks.dart';
import '../deferred_load/output_unit.dart';
import '../diagnostics/diagnostic_listener.dart';
import '../dump_info.dart';
import '../elements/entities.dart';
import '../environment.dart';
import '../inferrer/abstract_value_domain.dart';
import '../inferrer/abstract_value_strategy.dart';
import '../inferrer/types.dart';
import '../io/source_information.dart';
import '../js_backend/codegen_inputs.dart';
import '../js_backend/inferred_data.dart';
import '../js_model/js_strategy.dart';
import '../js_model/js_world.dart';
import '../js_model/locals.dart';
import '../options.dart';
import 'deferrable.dart';
import 'serialization.dart';

class _StringInterner implements ir.StringInterner, StringInterner {
  final Map<String, String> _map = {};

  @override
  String internString(String string) {
    return _map[string] ??= string;
  }
}

class SerializationTask extends CompilerTask {
  final CompilerOptions _options;
  final DiagnosticReporter _reporter;
  final api.CompilerInput _provider;
  final api.CompilerOutput _outputProvider;
  final _stringInterner = _StringInterner();
  final ValueInterner? _valueInterner;

  SerializationTask(
    this._options,
    this._reporter,
    this._provider,
    this._outputProvider,
    Measurer measurer,
  ) : _valueInterner = _options.features.internValues.isEnabled
          ? ValueInterner()
          : null,
      super(measurer);

  @override
  String get name => 'Serialization';

  void serializeComponent(
    ir.Component component, {
    bool includeSourceBytes = true,
  }) {
    measureSubtask('serialize dill', () {
      _reporter.log('Writing dill to ${_options.outputUri}');
      api.BinaryOutputSink dillOutput = _outputProvider.createBinarySink(
        _options.outputUri!,
      );
      ir.BinaryPrinter printer = ir.BinaryPrinter(
        dillOutput,
        includeSourceBytes: includeSourceBytes,
      );
      printer.writeComponentFile(component);
      dillOutput.close();
    });
  }

  Future<ir.Component> deserializeComponent() async {
    return measureIoSubtask('deserialize dill', () async {
      _reporter.log('Reading dill from ${_options.inputDillUri}');
      final dillInput = await _provider.readFromUri(
        _options.inputDillUri,
        inputKind: api.InputKind.binary,
      );
      ir.Component component = ir.Component();
      // Not using growable lists saves memory.
      ir.BinaryBuilder(
        dillInput.data,
        useGrowableLists: false,
        stringInterner: _stringInterner,
      ).readComponent(component);
      return component;
    });
  }

  Future<ir.Component> deserializeComponentAndUpdateOptions() async {
    ir.Component component = await deserializeComponent();
    return component;
  }

  void serializeClosedWorld(
    JClosedWorld closedWorld,
    SerializationIndices indices,
  ) {
    measureSubtask('serialize closed world', () {
      final outputUri = _options.dataUriForStage(CompilerStage.closedWorld);
      _reporter.log('Writing closed world to $outputUri');
      api.BinaryOutputSink dataOutput = _outputProvider.createBinarySink(
        outputUri,
      );
      DataSinkWriter sink = DataSinkWriter(
        BinaryDataSink(dataOutput),
        _options,
        indices,
      );
      serializeClosedWorldToSink(closedWorld, sink);
    });
  }

  Future<JClosedWorld> deserializeClosedWorld(
    AbstractValueStrategy abstractValueStrategy,
    ir.Component component,
    bool useDeferredSourceReads,
    SerializationIndices indices,
  ) async {
    return await measureIoSubtask('deserialize closed world', () async {
      final uri = _options.dataUriForStage(CompilerStage.closedWorld);
      _reporter.log('Reading data from $uri');
      api.Input<Uint8List> dataInput = await _provider.readFromUri(
        uri,
        inputKind: api.InputKind.binary,
      );
      DataSourceReader source = DataSourceReader(
        BinaryDataSource(dataInput.data, stringInterner: _stringInterner),
        _options,
        indices,
        interner: _valueInterner,
        useDeferredStrategy: useDeferredSourceReads,
      );
      var closedWorld = deserializeClosedWorldFromSource(
        _options,
        _reporter,
        abstractValueStrategy,
        component,
        source,
      );
      return closedWorld;
    });
  }

  void serializeGlobalTypeInference(
    GlobalTypeInferenceResults results,
    SerializationIndices indices,
  ) {
    measureSubtask('serialize data', () {
      final outputUri = _options.dataUriForStage(CompilerStage.globalInference);
      _reporter.log('Writing data to $outputUri');
      api.BinaryOutputSink dataOutput = _outputProvider.createBinarySink(
        outputUri,
      );
      DataSinkWriter sink = DataSinkWriter(
        BinaryDataSink(dataOutput),
        _options,
        indices,
      );
      serializeGlobalTypeInferenceResultsToSink(results, sink);
    });
  }

  Future<GlobalTypeInferenceResults> deserializeGlobalTypeInferenceResults(
    Environment environment,
    AbstractValueStrategy abstractValueStrategy,
    ir.Component component,
    JClosedWorld closedWorld,
    bool useDeferredSourceReads,
    SerializationIndices indices,
  ) async {
    return await measureIoSubtask('deserialize data', () async {
      final uri = _options.dataUriForStage(CompilerStage.globalInference);
      _reporter.log('Reading data from $uri');
      api.Input<Uint8List> dataInput = await _provider.readFromUri(
        uri,
        inputKind: api.InputKind.binary,
      );
      DataSourceReader source = DataSourceReader(
        BinaryDataSource(dataInput.data, stringInterner: _stringInterner),
        _options,
        indices,
        interner: _valueInterner,
        useDeferredStrategy: useDeferredSourceReads,
      );
      return deserializeGlobalTypeInferenceResultsFromSource(
        _options,
        _reporter,
        environment,
        abstractValueStrategy,
        component,
        closedWorld,
        source,
      );
    });
  }

  void serializeCodegen(
    JsBackendStrategy backendStrategy,
    AbstractValueDomain domain,
    CodegenResults codegenResults,
    SerializationIndices indices,
  ) {
    int shard = _options.codegenShard!;
    int shards = _options.codegenShards!;
    Map<MemberEntity, CodegenResult> results = {};
    int index = 0;
    final lazyMemberBodies = backendStrategy.forEachCodegenMember((
      MemberEntity member,
    ) {
      if (index % shards == shard) {
        final (result: codegenResult, isGenerated: _) = codegenResults
            .getCodegenResults(member);
        results[member] = codegenResult;
      }
      index++;
    });
    measureSubtask('serialize codegen', () {
      final outputUri = _options.dataUriForStage(CompilerStage.codegenSharded);
      Uri uri = Uri.parse('$outputUri$shard');
      api.BinaryOutputSink dataOutput = _outputProvider.createBinarySink(uri);
      DataSinkWriter sink = DataSinkWriter(
        BinaryDataSink(dataOutput),
        _options,
        indices,
      );
      _reporter.log('Writing data to $uri');
      sink.writeMembers(lazyMemberBodies);
      sink.registerAbstractValueDomain(domain);
      sink.writeMemberMap(results, (MemberEntity member, CodegenResult result) {
        sink.writeDeferrable(() => result.writeToDataSink(sink));
      });
      sink.close();
    });
  }

  Future<CodegenResults> deserializeCodegen(
    JsBackendStrategy backendStrategy,
    JClosedWorld closedWorld,
    CodegenInputs codegenInputs,
    bool useDeferredSourceReads,
    SourceLookup sourceLookup,
    SerializationIndices indices,
  ) async {
    int shards = _options.codegenShards!;
    Map<MemberEntity, Deferrable<CodegenResult>> results = {};
    for (int shard = 0; shard < shards; shard++) {
      Uri uri = Uri.parse(
        '${_options.dataUriForStage(CompilerStage.codegenSharded)}$shard',
      );
      await measureIoSubtask('deserialize codegen', () async {
        _reporter.log('Reading data from $uri');
        api.Input<Uint8List> dataInput = await _provider.readFromUri(
          uri,
          inputKind: api.InputKind.binary,
        );
        // TODO(36983): This code is extracted because there appeared to be a
        // memory leak for large buffer held by `source`.
        _deserializeCodegenInput(
          backendStrategy,
          closedWorld,
          uri,
          dataInput,
          results,
          useDeferredSourceReads,
          sourceLookup,
          indices,
        );
        dataInput.release();
      });
    }
    return DeserializedCodegenResults(
      codegenInputs,
      DeferrableValueMap(results),
      backendStrategy.functionCompiler,
    );
  }

  void _deserializeCodegenInput(
    JsBackendStrategy backendStrategy,
    JClosedWorld closedWorld,
    Uri uri,
    api.Input<Uint8List> dataInput,
    Map<MemberEntity, Deferrable<CodegenResult>> results,
    bool useDeferredSourceReads,
    SourceLookup sourceLookup,
    SerializationIndices indices,
  ) {
    DataSourceReader source = DataSourceReader(
      BinaryDataSource(dataInput.data, stringInterner: _stringInterner),
      _options,
      indices,
      interner: _valueInterner,
      useDeferredStrategy: useDeferredSourceReads,
    );
    backendStrategy.prepareCodegenReader(source);
    source.registerSourceLookup(sourceLookup);
    final lazyMemberBodies = source.readMembers();
    closedWorld.elementMap.registerLazyMemberBodies(lazyMemberBodies);
    source.registerAbstractValueDomain(closedWorld.abstractValueDomain);
    Map<MemberEntity, Deferrable<CodegenResult>> codegenResults = source
        .readMemberMap((MemberEntity member) {
          return source.readDeferrable(
            CodegenResult.readFromDataSource,
            cacheData: false,
          );
        });
    _reporter.log('Read ${codegenResults.length} members from $uri');
    results.addAll(codegenResults);
  }

  DataSinkWriter dataSinkWriterForDumpInfo(
    AbstractValueDomain abstractValueDomain,
    SerializationIndices indices,
  ) {
    final outputUri = _options.dataUriForStage(CompilerStage.dumpInfo);
    api.BinaryOutputSink dataOutput = _outputProvider.createBinarySink(
      outputUri,
    );
    final sink = DataSinkWriter(BinaryDataSink(dataOutput), _options, indices);
    sink.registerAbstractValueDomain(abstractValueDomain);
    return sink;
  }

  void serializeDumpInfoProgramData(
    DataSinkWriter sink,
    JsBackendStrategy backendStrategy,
    DumpInfoProgramData dumpInfoProgramData,
    DumpInfoJsAstRegistry dumpInfoRegistry,
  ) {
    dumpInfoProgramData.writeToDataSink(sink, dumpInfoRegistry);
    sink.close();
  }

  Future<DumpInfoProgramData> deserializeDumpInfoProgramData(
    JsBackendStrategy backendStrategy,
    AbstractValueDomain abstractValueDomain,
    OutputUnitData outputUnitData,
    SerializationIndices indices,
  ) async {
    final inputUri = _options.dataUriForStage(CompilerStage.dumpInfo);
    final dataInput = await _provider.readFromUri(
      inputUri,
      inputKind: api.InputKind.binary,
    );
    final source = DataSourceReader(
      BinaryDataSource(dataInput.data, stringInterner: _stringInterner),
      _options,
      indices,
      // This must use a deferred strategy so that we can delay reading the
      // registered impacts until we are able to read the count of them.
      useDeferredStrategy: true,
    );
    backendStrategy.prepareCodegenReader(source);
    source.registerAbstractValueDomain(abstractValueDomain);
    return DumpInfoProgramData.readFromDataSource(
      source,
      includeCodeText: _options.dumpInfoFormat != DumpInfoFormat.binary,
    );
  }
}

void serializeGlobalTypeInferenceResultsToSink(
  GlobalTypeInferenceResults results,
  DataSinkWriter sink,
) {
  final closedWorld = results.closedWorld;
  GlobalLocalsMap globalLocalsMap = results.globalLocalsMap;
  InferredData inferredData = results.inferredData;
  globalLocalsMap.writeToDataSink(sink);
  inferredData.writeToDataSink(sink);
  results.writeToDataSink(sink, closedWorld.elementMap);
  sink.close();
}

GlobalTypeInferenceResults deserializeGlobalTypeInferenceResultsFromSource(
  CompilerOptions options,
  DiagnosticReporter reporter,
  Environment environment,
  AbstractValueStrategy abstractValueStrategy,
  ir.Component component,
  JClosedWorld closedWorld,
  DataSourceReader source,
) {
  source.registerComponentLookup(ComponentLookup(component));
  GlobalLocalsMap globalLocalsMap = GlobalLocalsMap.readFromDataSource(
    closedWorld.closureDataLookup.getEnclosingMember,
    source,
  );
  InferredData inferredData = InferredData.readFromDataSource(
    source,
    closedWorld,
  );
  return GlobalTypeInferenceResults.readFromDataSource(
    source,
    closedWorld.elementMap,
    closedWorld,
    globalLocalsMap,
    inferredData,
  );
}

void serializeClosedWorldToSink(JClosedWorld closedWorld, DataSinkWriter sink) {
  closedWorld.writeToDataSink(sink);
  sink.close();
}

JClosedWorld deserializeClosedWorldFromSource(
  CompilerOptions options,
  DiagnosticReporter reporter,
  AbstractValueStrategy abstractValueStrategy,
  ir.Component component,
  DataSourceReader source,
) {
  return JClosedWorld.readFromDataSource(
    options,
    reporter,
    abstractValueStrategy,
    component,
    source,
  );
}