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

import Fuzzilli

fileprivate let ForceJITCompilationThroughLoopGenerator = CodeGenerator("ForceJITCompilationThroughLoopGenerator", inputs: .required(.function())) { b, f in
    assert(b.type(of: f).Is(.function()))
    let arguments = b.randomArguments(forCalling: f)

    b.buildRepeatLoop(n: 100) { _ in
        b.callFunction(f, withArgs: arguments)
    }
}

fileprivate let ForceTurboFanCompilationGenerator = CodeGenerator("ForceTurboFanCompilationGenerator", inputs: .required(.function())) { b, f in
    assert(b.type(of: f).Is(.function()))
    let arguments = b.randomArguments(forCalling: f)

    b.callFunction(f, withArgs: arguments)

    b.eval("%PrepareFunctionForOptimization(%@)", with: [f]);

    b.callFunction(f, withArgs: arguments)
    b.callFunction(f, withArgs: arguments)

    b.eval("%OptimizeFunctionOnNextCall(%@)", with: [f]);

    b.callFunction(f, withArgs: arguments)
}

fileprivate let ForceMaglevCompilationGenerator = CodeGenerator("ForceMaglevCompilationGenerator", inputs: .required(.function())) { b, f in
    assert(b.type(of: f).Is(.function()))
    let arguments = b.randomArguments(forCalling: f)

    b.callFunction(f, withArgs: arguments)

    b.eval("%PrepareFunctionForOptimization(%@)", with: [f]);

    b.callFunction(f, withArgs: arguments)
    b.callFunction(f, withArgs: arguments)

    b.eval("%OptimizeMaglevOnNextCall(%@)", with: [f]);

    b.callFunction(f, withArgs: arguments)
}

fileprivate let TurbofanVerifyTypeGenerator = CodeGenerator("TurbofanVerifyTypeGenerator", inputs: .one) { b, v in
    b.eval("%VerifyType(%@)", with: [v])
}

// Insert random GC calls throughout our code.
fileprivate let GcGenerator = CodeGenerator("GcGenerator") { b in
    let gc = b.createNamedVariable(forBuiltin: "gc")

    // Do minor GCs more frequently.
    let type = b.loadString(probability(0.25) ? "major" : "minor")
    // If the execution type is 'async', gc() returns a Promise, we currently
    // do not really handle other than typing the return of gc to .undefined |
    // .jsPromise. One could either chain a .then or create two wrapper
    // functions that are differently typed such that fuzzilli always knows
    // what the type of the return value is.
    let execution = b.loadString(probability(0.5) ? "sync" : "async")
    b.callFunction(gc, withArgs: [b.createObject(with: ["type": type, "execution": execution])])
}

public extension ILType {
{% for enum in database.enumerations %}
  static let js{{enum.identifier}} = ILType.enumeration(ofName: "{{enum.identifier}}", withValues: [ {{enum.values|map("tojson")|join(", ")}} ])
{% endfor %}
{% for interface in database.interfaces|sort_object_groups %}
  {% set inter = interface|parse_interface -%}
  static let js{{interface.identifier}} = {{inter[0].fuzzilli_repr()}}
{% endfor %}
{% for d in database.dictionaries|sort_object_groups %}
  {% set dictionary = d|parse_dictionary -%}
  static let js{{d.identifier}} = {{dictionary[0].fuzzilli_repr()}}
{% endfor %}
{% for interface in database.interfaces|sort_object_groups -%}
  {% set constructor = interface|parse_constructors -%}
{%   if constructor[0] %}
  static let js{{interface.identifier}}Constructor = {{constructor[0].fuzzilli_repr()}}
{%   endif %}
{%- endfor %}
{% for t in database.typedefs %}
  {% set il_type = t.idl_type|idl_type_to_iltype -%}
  static let js{{t.identifier}} = {{il_type.fuzzilli_repr()}}
{% endfor %}
{% for interface in database.callback_interfaces %}
  {% set inter = interface|parse_interface -%}
  static let js{{interface.identifier}} = {{inter[0].fuzzilli_repr()}}
{% endfor %}
{% for c in database.callback_functions %}
  {% set op = c|parse_operation -%}
  static let js{{c.identifier}} = ILType.function({{op.fuzzilli_repr()}})
{% endfor %}
}

{%- macro define_group(identifier, group) %}
fileprivate let js{{identifier}} = ObjectGroup(
  name: "{{group.name}}",
  instanceType: {{group.instanceType.fuzzilli_repr()}},
  properties: [
  {%- for k, v in group.properties.items() %}
    "{{k}}" : {{v.fuzzilli_repr()}},
  {% else %}
    :
  {% endfor -%}
  ],
  methods: [
  {%- for k, v in group.methods.items() %}
    "{{k}}" : {{v.fuzzilli_repr()}},
  {% else %}
    :
  {% endfor -%}
  ]{%- if group.parent %},
  parent: "{{group.parent}}"
  {%- endif %}
)
{%- endmacro -%}


{% for interface in database.interfaces|sort_object_groups %}
  {%- set inter = interface|parse_interface -%}
{{define_group(interface.identifier, inter[1])}}
{% endfor %}
{% for interface in database.interfaces|sort_object_groups %}
  {%- set constructor = interface|parse_constructors -%}
{%   if constructor[1] %}
{{define_group(interface.identifier+"Constructor", constructor[1])}}
{%   endif %}
{%- endfor %}
{% for d in database.dictionaries|sort_object_groups %}
  {%- set dictionary = d|parse_dictionary -%}
{{define_group(d.identifier, dictionary[1])}}
{% endfor -%}
{% for interface in database.callback_interfaces %}
  {%- set inter = interface|parse_interface -%}
{{define_group(interface.identifier, inter[1])}}
{% endfor -%}

fileprivate let fastCallables : [(group: ILType, method: String)] = [
{%- for interface in database.interfaces -%}
{%- for operation in interface.operations -%}
{%- if operation.is_static == False -%}
{%- if 'NoAllocDirectCall' in operation.extended_attributes %}
  (group: js{{interface.identifier}}.instanceType, method: "{{operation.identifier}}"),
{%- endif %}
{%- endif -%}
{%- endfor -%}
{%- endfor %}
]

fileprivate let FastApiCallFuzzer = ProgramTemplate("FastApiCallFuzzer") { b in
    b.buildPrefix()
    b.build(n: 20)
    let parameterCount = probability(0.5) ? 0 : Int.random(in: 1...4)

    let f = b.buildPlainFunction(with: .parameters(n: parameterCount)) { args in
        b.build(n: 10)
        let target = fastCallables.randomElement()!
        let apiObj = b.findOrGenerateType(target.group)
        let functionSig = chooseUniform(from: b.methodSignatures(of: target.method, on: target.group))
        let apiCall = b.callMethod(target.method, on: apiObj, withArgs: b.findOrGenerateArguments(forSignature: functionSig), guard: true)
        b.doReturn(apiCall)
    }

    let args = b.randomJsVariables(n: Int.random(in: 0...5))
    b.callFunction(f, withArgs: args)

    b.eval("%PrepareFunctionForOptimization(%@)", with: [f]);

    b.callFunction(f, withArgs: args)
    b.callFunction(f, withArgs: args)

    b.eval("%OptimizeFunctionOnNextCall(%@)", with: [f]);

    b.callFunction(f, withArgs: args)

    b.build(n: 10)
}

fileprivate let WasmFastCallFuzzer = WasmProgramTemplate("WasmFastCallFuzzer") { b in
    b.buildPrefix()
    b.build(n: 10)
    let target = fastCallables.randomElement()!
    let apiObj = b.findOrGenerateType(target.group)

    // Bind the API function so that it can be called from WebAssembly.
    let wrapped = b.bindMethod(target.method, on: apiObj)

    let functionSig = chooseUniform(from: b.methodSignatures(of: target.method, on: target.group))
    let wrappedSig = Signature(expects: [.plain(b.type(of: apiObj))] + functionSig.parameters, returns: functionSig.outputType)

    let m = b.buildWasmModule { m in
        let allWasmTypes: WeightedList<ILType> = WeightedList([(.wasmi32, 1), (.wasmi64, 1), (.wasmf32, 1), (.wasmf64, 1), (.wasmExternRef, 1), (.wasmFuncRef, 1)])
        let wasmSignature = ProgramBuilder.convertJsSignatureToWasmSignature(wrappedSig, availableTypes: allWasmTypes)
        m.addWasmFunction(with: wasmSignature) {fbuilder, _, _  in
            let args = b.randomWasmArguments(forWasmSignature: wasmSignature)
            if let args {
                let maybeRet = fbuilder.wasmJsCall(function: wrapped, withArgs: args, withWasmSignature: wasmSignature)
                if let ret = maybeRet {
                  return [ret]
                }
            } else {
                logger.error("Arguments should have been generated")
            }
            return wasmSignature.outputTypes.map(fbuilder.findOrGenerateWasmVar)
        }
    }

    let exports = m.loadExports()

    for (methodName, _) in m.getExportedMethods() {
        let exportedMethod = b.getProperty(methodName, of: exports)
        b.eval("%WasmTierUpFunction(%@)", with: [exportedMethod])
        let args = b.findOrGenerateArguments(forSignature: wrappedSig)
        b.callMethod(methodName, on: exports, withArgs: args)
    }
}



let chromiumProfile = Profile(
    processArgs: { _ in
        var args: [String] = []
        return args
    },

    processEnv: ["ASAN_OPTIONS":"detect_odr_violation=0", "DISPLAY":":20"],

    maxExecsBeforeRespawn: 100,

    timeout: 8000,

    codePrefix: """
                """,

    codeSuffix: """
                """,

    ecmaVersion: ECMAScriptVersion.es6,

    startupTests: [
        ("fuzzilli('FUZZILLI_CRASH', 0)", .shouldCrash),
        ("fuzzilli('FUZZILLI_CRASH', 1)", .shouldCrash),
        ("fuzzilli('FUZZILLI_CRASH', 2)", .shouldCrash),
        ("fuzzilli('FUZZILLI_CRASH', 3)", .shouldCrash),
    ],

    additionalCodeGenerators: [
        (ForceJITCompilationThroughLoopGenerator,  5),
        (ForceTurboFanCompilationGenerator,        5),
        (ForceMaglevCompilationGenerator,          5),
        (TurbofanVerifyTypeGenerator,             10),
        (GcGenerator,                             10),
    ],

    additionalProgramTemplates: WeightedList<ProgramTemplate>([
        (FastApiCallFuzzer,      5),
        (WasmFastCallFuzzer,     5),
    ]),

    disabledCodeGenerators: [],

    disabledMutators: [],

    additionalBuiltins: [
{%- for interface in database.interfaces|sort_object_groups -%}
  {% set constructor = interface|parse_constructors -%}
{%   if constructor[0] %}
      "{{interface.identifier}}": ILType.js{{interface.identifier}}Constructor,
{%   endif %}
{%- endfor %}
      "window": .jsWindow,
      "document": .jsDocument,
      "gc" : .function([] => (.undefined | .jsPromise)),
      "d8": .object(),
    ],

    additionalObjectGroups: [
{%- for interface in database.interfaces|sort_object_groups %}
      js{{interface.identifier}},
{% endfor -%}
{% for interface in database.interfaces|sort_object_groups -%}
  {% set constructor = interface|parse_constructors -%}
{%   if constructor[1] %}
      js{{interface.identifier}}Constructor,
{%   endif %}
{%- endfor -%}
{% for d in database.dictionaries|sort_object_groups %}
      js{{d.identifier}},
{% endfor -%}
{% for interface in database.callback_interfaces %}
      js{{interface.identifier}},
{% endfor -%}
    ],

    optionalPostProcessor: nil
)