import _ from 'lodash';
import '../../../../assert/js/assert.js';
import '../../../../globals.js';
// Copyright 2023-2024, University of Colorado Boulder
/**
 * TODO: doc
 *
 * @author Jonathan Olson <jonathan.olson@colorado.edu>
 */
import { alpenglow, BindGroupLayout, BindingDescriptor, BufferSlot, BufferSlotSlice, PipelineLayout } from '../../imports.js';
class Routine {
    constructor(deviceContext, module, nonBufferSlots, rootBufferSlots, bufferSliceMap, pipelineLayoutMap, computePipelineMap, 
    // TODO: factor out some types?
    executeWrapper) {
        this.deviceContext = deviceContext;
        this.module = module;
        this.nonBufferSlots = nonBufferSlots;
        this.rootBufferSlots = rootBufferSlots;
        this.bufferSliceMap = bufferSliceMap;
        this.pipelineLayoutMap = pipelineLayoutMap;
        this.computePipelineMap = computePipelineMap;
        this.executeWrapper = executeWrapper;
        this.pipelineBlueprints = module.pipelineBlueprints;
        this.rootResourceSlots = [...nonBufferSlots, ...rootBufferSlots];
        this.bindGroupLayouts = _.uniq([...this.pipelineLayoutMap.values()].flatMap(layout => layout.bindGroupLayouts));
    }
    execute(context, data) {
        return this.executeWrapper(context, this.module.execute, data);
    }
    static async create(deviceContext, module, sharedBufferSlots, layoutStrategy, executeWrapper) {
        const slots = _.uniq([
            ...sharedBufferSlots,
            ...module.getResourceSlots()
        ]);
        const nonBufferSlots = slots.filter(slot => !(slot instanceof BufferSlot));
        const bufferSlots = slots.filter(slot => slot instanceof BufferSlot);
        const rootBufferSlots = bufferSlots.filter(slot => {
            return !bufferSlots.some(otherSlot => otherSlot.hasChildSlot(slot));
        });
        const bufferSliceMap = new Map();
        const recur = (rootSlot, slot, offset) => {
            bufferSliceMap.set(slot, new BufferSlotSlice(rootSlot, offset));
            for (const slice of slot.bufferSlotSlices) {
                recur(rootSlot, slice.bufferSlot, offset + slice.offset);
            }
        };
        for (const slot of rootBufferSlots) {
            recur(slot, slot, 0);
        }
        // NOTE: Do bind group/pipeline layouts AFTER we figure out buffer slices recursively, since if we add dynamic
        // offsets, we'll need to know that before computing the layouts.
        const pipelineLayoutMap = layoutStrategy(deviceContext, module.pipelineBlueprints);
        const computePipelineMap = new Map();
        for (const pipelineBlueprint of _.uniq(module.pipelineBlueprints)) {
            const pipelineLayout = pipelineLayoutMap.get(pipelineBlueprint);
            assert && assert(pipelineLayout, 'Missing pipeline layout');
            computePipelineMap.set(pipelineBlueprint, await pipelineBlueprint.toComputePipeline(deviceContext, pipelineLayout));
        }
        return new Routine(deviceContext, module, nonBufferSlots, rootBufferSlots, bufferSliceMap, pipelineLayoutMap, computePipelineMap, executeWrapper);
    }
}
// TODO: Since we're returning a map, it's enforcing "no different layouts for different locations of the same blueprint"
// TODO: which is less than ideal(!). Figure this out, and possibly different bind groups for the same bind group layout.
Routine.INDIVIDUAL_LAYOUT_STRATEGY = (deviceContext, pipelineBlueprints) => {
    const map = new Map();
    _.uniq(pipelineBlueprints).forEach(pipelineBlueprint => {
        const bindGroupLayout = new BindGroupLayout(deviceContext, pipelineBlueprint.name, 0, pipelineBlueprint.usages.map((usage, index) => {
            return new BindingDescriptor(index, usage.bindingType, usage.resourceSlot);
        }));
        const pipelineLayout = new PipelineLayout(deviceContext, [bindGroupLayout]);
        map.set(pipelineBlueprint, pipelineLayout);
    });
    return map;
};
Routine.COMBINE_ALL_LAYOUT_STRATEGY = (deviceContext, pipelineBlueprints) => {
    const slots = _.uniq(pipelineBlueprints.flatMap(pipelineBlueprint => pipelineBlueprint.usages.map(usage => usage.resourceSlot)));
    const bindingDescriptors = slots.map((slot, index) => {
        // Compute the binding type
        let bindingType = null;
        for (const pipelineBlueprint of pipelineBlueprints) {
            const usage = pipelineBlueprint.usages.find(usage => usage.resourceSlot === slot);
            if (usage) {
                if (!bindingType) {
                    bindingType = usage.bindingType;
                }
                else {
                    bindingType = bindingType.combined(usage.bindingType);
                }
            }
        }
        assert && assert(bindingType, 'Missing binding type');
        return new BindingDescriptor(index, bindingType, slot);
    });
    const bindGroupLayout = new BindGroupLayout(deviceContext, 'combine all', 0, bindingDescriptors);
    const pipelineLayout = new PipelineLayout(deviceContext, [bindGroupLayout]);
    const map = new Map();
    _.uniq(pipelineBlueprints).forEach(pipelineBlueprint => {
        map.set(pipelineBlueprint, pipelineLayout);
    });
    return map;
};
export default Routine;
alpenglow.register('Routine', Routine);
//# sourceMappingURL=Routine.js.map