// Copyright (c) 2022, 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 'package:kernel/ast.dart';
import 'package:wasm_builder/wasm_builder.dart' as w;

import 'class_info.dart';
import 'closures.dart';
import 'dispatch_table.dart';
import 'reference_extensions.dart';
import 'translator.dart';

/// This class is responsible for collecting import and export annotations.
/// It also creates Wasm functions for Dart members and manages the worklist
/// used to achieve tree shaking.
class FunctionCollector {
  final Translator translator;

  // Wasm function for each Dart function
  final Map<Reference, w.BaseFunction> _functions = {};
  // Names of exported functions
  final Map<Reference, String> _exports = {};
  // Functions for which code has not yet been generated
  final List<Reference> _worklist = [];
  // Selector IDs that are invoked via GDT.
  final Set<int> _calledSelectors = {};
  // Class IDs for classes that are allocated somewhere in the program
  final Set<int> _allocatedClasses = {};
  // For each class ID, which functions should be added to the worklist if an
  // allocation of that class is encountered
  final Map<int, List<Reference>> _pendingAllocation = {};

  FunctionCollector(this.translator);

  w.ModuleBuilder get m => translator.m;

  void collectImportsAndExports() {
    for (Library library in translator.libraries) {
      library.procedures.forEach(_importOrExport);
      library.fields.forEach(_importOrExport);
      for (Class cls in library.classes) {
        cls.procedures.forEach(_importOrExport);
      }
    }
  }

  bool isWorkListEmpty() => _worklist.isEmpty;

  Reference popWorkList() => _worklist.removeLast();

  void _importOrExport(Member member) {
    String? importName =
        translator.getPragma(member, "wasm:import", member.name.text);
    if (importName != null) {
      int dot = importName.indexOf('.');
      if (dot != -1) {
        assert(!member.isInstanceMember);
        String module = importName.substring(0, dot);
        String name = importName.substring(dot + 1);
        if (member is Procedure) {
          // Define the function type in a singular recursion group to enable it
          // to be unified with function types defined in FFI modules or using
          // `WebAssembly.Function`.
          m.types.splitRecursionGroup();
          w.FunctionType ftype = _makeFunctionType(
              translator, member.reference, null,
              isImportOrExport: true);
          m.types.splitRecursionGroup();
          _functions[member.reference] =
              m.functions.import(module, name, ftype, "$importName (import)");
        }
      }
    }
    String? exportName =
        translator.getPragma(member, "wasm:export", member.name.text);
    if (exportName != null) {
      if (member is Procedure) {
        // Although we don't need type unification for the types of exported
        // functions, we still place these types in singleton recursion groups,
        // since Binaryen's `--closed-world` optimization mode requires all
        // publicly exposed types to be defined in separate recursion groups
        // from GC types.
        m.types.splitRecursionGroup();
        _makeFunctionType(translator, member.reference, null,
            isImportOrExport: true);
        m.types.splitRecursionGroup();
      }
      addExport(member.reference, exportName);
    }
  }

  void addExport(Reference target, String exportName) {
    _exports[target] = exportName;
  }

  String? getExport(Reference target) => _exports[target];

  void initialize() {
    // Add exports to the module and add exported functions to the worklist
    for (var export in _exports.entries) {
      Reference target = export.key;
      Member node = target.asMember;
      if (node is Procedure) {
        _worklist.add(target);
        assert(!node.isInstanceMember);
        assert(!node.isGetter);
        w.FunctionType ftype =
            _makeFunctionType(translator, target, null, isImportOrExport: true);
        w.BaseFunction function = m.functions.define(ftype, "$node");
        _functions[target] = function;
        m.exports.export(export.value, function);
      } else if (node is Field) {
        w.Table? table = translator.getTable(node);
        if (table != null) {
          m.exports.export(export.value, table);
        }
      }
    }

    // Value classes are always implicitly allocated.
    recordClassAllocation(
        translator.classInfo[translator.boxedBoolClass]!.classId);
    recordClassAllocation(
        translator.classInfo[translator.boxedIntClass]!.classId);
    recordClassAllocation(
        translator.classInfo[translator.boxedDoubleClass]!.classId);
  }

  w.BaseFunction? getExistingFunction(Reference target) {
    return _functions[target];
  }

  w.BaseFunction getFunction(Reference target) {
    return _functions.putIfAbsent(target, () {
      _worklist.add(target);
      return _getFunctionTypeAndName(target, m.functions.define);
    });
  }

  w.FunctionType getFunctionType(Reference target) {
    return _getFunctionTypeAndName(target, (ftype, name) => ftype);
  }

  /// Pass the Wasm type and name of the function for [target] to [action].
  ///
  /// Name should be used for the Wasm names section entry for the function so
  /// that the error stack traces will have names expected by the Dart spec.
  T _getFunctionTypeAndName<T>(
      Reference target, T Function(w.FunctionType, String) action) {
    if (target.isTypeCheckerReference) {
      Member member = target.asMember;
      if (member is Field || (member is Procedure && member.isSetter)) {
        return action(translator.dynamicSetForwarderFunctionType,
            '${target.asMember} setter type checker');
      } else {
        return action(translator.dynamicInvocationForwarderFunctionType,
            '${target.asMember} invocation type checker');
      }
    }

    if (target.isTearOffReference) {
      return action(
          translator.dispatchTable.selectorForTarget(target).signature,
          "${target.asMember} tear-off");
    }

    Member member = target.asMember;
    final ftype = member.accept1(_FunctionTypeGenerator(translator), target);

    String memberName = member.toString();
    if (memberName.endsWith('.')) {
      memberName = memberName.substring(0, memberName.length - 1);
    }

    if (target.isInitializerReference) {
      return action(ftype, 'new $memberName (initializer)');
    } else if (target.isConstructorBodyReference) {
      return action(ftype, 'new $memberName (constructor body)');
    } else if (member is Procedure && member.isFactory) {
      return action(ftype, 'new $memberName');
    } else {
      return action(ftype, memberName);
    }
  }

  void recordSelectorUse(SelectorInfo selector) {
    if (_calledSelectors.add(selector.id)) {
      selector.targets.forEach((classId, target) {
        if (!target.asMember.isAbstract) {
          if (_allocatedClasses.contains(classId)) {
            // Class declaring or inheriting member is allocated somewhere.
            getFunction(target);
          } else {
            // Remember the member in case an allocation is encountered later.
            _pendingAllocation.putIfAbsent(classId, () => []).add(target);
          }
        }
      });
    }
  }

  void recordClassAllocation(int classId) {
    if (_allocatedClasses.add(classId)) {
      // Schedule all members that were pending allocation of this class.
      for (Reference target in _pendingAllocation[classId] ?? const []) {
        getFunction(target);
      }
    }
  }

  /// Returns an iterable of translated procedures.
  Iterable<Procedure> get translatedProcedures =>
      _functions.keys.map((k) => k.node).whereType<Procedure>();
}

class _FunctionTypeGenerator extends MemberVisitor1<w.FunctionType, Reference> {
  final Translator translator;

  _FunctionTypeGenerator(this.translator);

  @override
  w.FunctionType visitField(Field node, Reference target) {
    if (!node.isInstanceMember) {
      if (target == node.fieldReference) {
        // Static field initializer function
        return _makeFunctionType(translator, target, null);
      }
      String kind = target == node.setterReference ? "setter" : "getter";
      throw "No implicit $kind function for static field: $node";
    }
    return translator.dispatchTable.selectorForTarget(target).signature;
  }

  @override
  w.FunctionType visitProcedure(Procedure node, Reference target) {
    assert(!node.isAbstract);
    return node.isInstanceMember
        ? translator.dispatchTable.selectorForTarget(node.reference).signature
        : _makeFunctionType(translator, target, null);
  }

  @override
  w.FunctionType visitConstructor(Constructor node, Reference target) {
    // Get this constructor's argument types
    List<w.ValueType> arguments = _getInputTypes(
        translator, target, null, false, translator.translateType);

    if (translator.constructorClosures[node.reference] == null) {
      // We need the contexts of the constructor before generating the
      // initializer and constructor body functions, as these functions will
      // return/take a context argument if context must be shared between them.
      // Generate the contexts the first time we visit a constructor.
      Closures closures = Closures(translator, node);

      closures.findCaptures(node);
      closures.collectContexts(node);
      closures.buildContexts();

      translator.constructorClosures[node.reference] = closures;
    }

    if (target.isInitializerReference) {
      return _getInitializerType(node, target, arguments);
    }

    if (target.isConstructorBodyReference) {
      return _getConstructorBodyType(node, arguments);
    }

    return _getConstructorAllocatorType(node, arguments);
  }

  w.FunctionType _getConstructorAllocatorType(
      Constructor node, List<w.ValueType> arguments) {
    return translator.m.types.defineFunction(arguments,
        [translator.classInfo[node.enclosingClass]!.nonNullableType.unpacked]);
  }

  w.FunctionType _getInitializerType(
      Constructor node, Reference target, List<w.ValueType> arguments) {
    final ClassInfo info = translator.classInfo[node.enclosingClass]!;
    assert(translator.constructorClosures.containsKey(node.reference));
    Closures closures = translator.constructorClosures[node.reference]!;

    List<w.ValueType> superOrRedirectedInitializerArgs = [];

    for (Initializer initializer in node.initializers) {
      if (initializer is SuperInitializer) {
        Supertype? supersupertype = initializer.target.enclosingClass.supertype;

        if (supersupertype != null) {
          ClassInfo superInfo = info.superInfo!;
          w.FunctionType superInitializer = translator.functions
              .getFunctionType(initializer.target.initializerReference);

          final int numSuperclassFields = superInfo.getClassFieldTypes().length;
          final int numSuperContextAndConstructorArgs =
              superInitializer.outputs.length - numSuperclassFields;

          // get types of super initializer outputs, ignoring the superclass
          // fields
          superOrRedirectedInitializerArgs = superInitializer.outputs
              .sublist(0, numSuperContextAndConstructorArgs);
        }
      } else if (initializer is RedirectingInitializer) {
        Supertype? supersupertype = initializer.target.enclosingClass.supertype;

        if (supersupertype != null) {
          w.FunctionType redirectedInitializer = translator.functions
              .getFunctionType(initializer.target.initializerReference);

          final int numClassFields = info.getClassFieldTypes().length;
          final int numRedirectedContextAndConstructorArgs =
              redirectedInitializer.outputs.length - numClassFields;

          // get types of redirecting initializer outputs, ignoring the class
          // fields
          superOrRedirectedInitializerArgs = redirectedInitializer.outputs
              .sublist(0, numRedirectedContextAndConstructorArgs);
        }
      }
    }

    // Get this classes's field types
    final List<w.ValueType> fieldTypes = info.getClassFieldTypes();

    // Add nullable context reference for when the constructor has a non-empty
    // context
    Context? context = closures.contexts[node];
    w.ValueType? contextRef;

    if (context != null) {
      assert(!context.isEmpty);
      contextRef = w.RefType.struct(nullable: true);
    }

    final List<w.ValueType> outputs = superOrRedirectedInitializerArgs +
        arguments.reversed.toList() +
        (contextRef != null ? [contextRef] : []) +
        fieldTypes;

    return translator.m.types.defineFunction(arguments, outputs);
  }

  w.FunctionType _getConstructorBodyType(
      Constructor node, List<w.ValueType> arguments) {
    assert(translator.constructorClosures.containsKey(node.reference));
    Closures closures = translator.constructorClosures[node.reference]!;
    Context? context = closures.contexts[node];

    List<w.ValueType> inputs = [
      translator.classInfo[node.enclosingClass]!.nonNullableType.unpacked
    ];

    if (context != null) {
      assert(!context.isEmpty);
      // Nullable context reference for when the constructor has a non-empty
      // context
      w.ValueType contextRef = w.RefType.struct(nullable: true);
      inputs.add(contextRef);
    }

    inputs += arguments;

    for (Initializer initializer in node.initializers) {
      if (initializer is SuperInitializer ||
          initializer is RedirectingInitializer) {
        Constructor target = initializer is SuperInitializer
            ? initializer.target
            : (initializer as RedirectingInitializer).target;

        Supertype? supersupertype = target.enclosingClass.supertype;

        if (supersupertype != null) {
          w.FunctionType superOrRedirectedConstructorBodyType = translator
              .functions
              .getFunctionType(target.constructorBodyReference);

          // drop receiver param
          inputs += superOrRedirectedConstructorBodyType.inputs.sublist(1);
        }
      }
    }

    return translator.m.types.defineFunction(inputs, []);
  }
}

List<w.ValueType> _getInputTypes(
    Translator translator,
    Reference target,
    w.ValueType? receiverType,
    bool isImportOrExport,
    w.ValueType Function(DartType) translateType) {
  Member member = target.asMember;
  int typeParamCount = 0;
  Iterable<DartType> params;
  if (member is Field) {
    params = [if (target.isImplicitSetter) member.setterType];
  } else {
    FunctionNode function = member.function!;
    typeParamCount = (member is Constructor
            ? member.enclosingClass.typeParameters
            : function.typeParameters)
        .length;
    List<String> names = [for (var p in function.namedParameters) p.name!]
      ..sort();
    final typeForParam = translator.typeOfParameterVariable;
    Map<String, DartType> nameTypes = {
      for (var p in function.namedParameters)
        p.name!: typeForParam(p, p.isRequired)
    };
    final positionals = function.positionalParameters;
    params = [
      for (int i = 0; i < positionals.length; ++i)
        typeForParam(positionals[i], i < function.requiredParameterCount),
      for (String name in names) nameTypes[name]!
    ];
  }

  final List<w.ValueType> typeParameters = List.filled(
      typeParamCount,
      translateType(
          InterfaceType(translator.typeClass, Nullability.nonNullable)));

  final List<w.ValueType> inputs = [];

  if (receiverType != null) {
    assert(!isImportOrExport);
    inputs.add(receiverType);
  }

  inputs.addAll(typeParameters);
  inputs.addAll(params.map(translateType));

  return inputs;
}

w.FunctionType _makeFunctionType(
    Translator translator, Reference target, w.ValueType? receiverType,
    {bool isImportOrExport = false}) {
  Member member = target.asMember;

  // Translate types differently for imports and exports.
  w.ValueType translateType(DartType type) => isImportOrExport
      ? translator.translateExternalType(type)
      : translator.translateType(type);

  final List<w.ValueType> inputs = _getInputTypes(
      translator, target, receiverType, isImportOrExport, translateType);

  // Mutable fields have initializer setters with a non-empty output list,
  // so check that the member is a Procedure
  final bool emptyOutputList = member is Procedure && member.isSetter;

  bool isVoidType(DartType t) =>
      (isImportOrExport && t is VoidType) ||
      (t is InterfaceType && t.classNode == translator.wasmVoidClass);

  final List<w.ValueType> outputs;
  if (emptyOutputList) {
    outputs = const [];
  } else {
    final DartType returnType = translator.typeOfReturnValue(member);
    outputs = !isVoidType(returnType) ? [translateType(returnType)] : const [];
  }

  return translator.m.types.defineFunction(inputs, outputs);
}