luau-bytecode 0.732.0

Luau bytecode model, builder, serializer, and dumper
Documentation
use super::*;

impl FunctionGraphBuilder {
    fn make_phi(&mut self, block: BytecodeBlockId, register: Register) -> BytecodeOperand {
        self.phis.push(BytecodePhi {
            operands: Vec::new(),
            users: Vec::new(),
        });
        let id = BytecodePhiId::new(self.phis.len() - 1);
        self.block_phis[block.index()].push(id);
        let operand = BytecodeOperand::Phi(id);
        self.registers.insert(operand, register);
        operand
    }

    pub(super) fn producer(
        &mut self,
        blocks: &[BytecodeBlock],
        register: Register,
    ) -> Result<BytecodeOperand, BytecodeReadError> {
        if let Some(operand) = self.read_variable(blocks, self.current_block, register) {
            return Ok(operand);
        }

        if self.is_unreachable(blocks, self.current_block, &mut HashSet::new()) {
            return Ok(BytecodeOperand::VmRegister(register));
        }

        Err(BytecodeReadError::MissingRegisterProducer {
            block: self.current_block.index(),
            register,
        })
    }

    fn read_variable(
        &mut self,
        blocks: &[BytecodeBlock],
        block: BytecodeBlockId,
        register: Register,
    ) -> Option<BytecodeOperand> {
        let block_producers = &self.producers.blocks[block.index()];
        if i32::from(register) > block_producers.invalid_after {
            return None;
        }
        if let Some(operand) = block_producers.own.get(&register).copied() {
            return Some(operand);
        }
        if let Some(operand) = block_producers.cached.get(&register).copied() {
            return Some(operand);
        }
        if let Some(multi_return) = block_producers.multi_return
            && register >= block_producers.multi_return_start
        {
            let projection = self.projection(
                multi_return,
                u32::from(register - block_producers.multi_return_start),
            );
            self.producers.blocks[block.index()]
                .cached
                .insert(register, projection);
            return Some(projection);
        }

        Some(self.read_variable_recursive(blocks, block, register))
    }

    fn read_variable_recursive(
        &mut self,
        blocks: &[BytecodeBlock],
        block: BytecodeBlockId,
        register: Register,
    ) -> BytecodeOperand {
        if !self.producers.blocks[block.index()].sealed {
            let phi = self.make_phi(block, register);
            let producers = &mut self.producers.blocks[block.index()];
            producers.incomplete_phis.insert(register, phi);
            producers.cached.insert(register, phi);
            return phi;
        }

        let predecessors = blocks[block.index()].predecessors();
        if predecessors.is_empty() {
            return BytecodeOperand::VmRegister(register);
        }

        if predecessors.len() == 1 {
            let undefined = BytecodeOperand::VmRegister(register);
            self.producers.blocks[block.index()]
                .cached
                .insert(register, undefined);
            let value = self
                .read_variable(blocks, predecessors[0].target, register)
                .unwrap_or(undefined);
            self.producers.blocks[block.index()]
                .cached
                .insert(register, value);
            return value;
        }

        let phi = self.make_phi(block, register);
        self.producers.blocks[block.index()]
            .cached
            .insert(register, phi);
        self.add_phi_operands(blocks, block, register, phi);
        phi
    }

    fn add_phi_operands(
        &mut self,
        blocks: &[BytecodeBlock],
        block: BytecodeBlockId,
        register: Register,
        phi: BytecodeOperand,
    ) {
        let BytecodeOperand::Phi(phi_id) = phi else {
            unreachable!("phi operand must contain a phi id");
        };
        let predecessors = blocks[block.index()].predecessors().to_vec();
        for predecessor in predecessors {
            if let Some(operand) = self.read_variable(blocks, predecessor.target, register) {
                self.phis[phi_id.index()].operands.push(operand);
            }
        }
    }

    pub(in crate::graph) fn finish_block(
        &mut self,
        block: BytecodeBlockId,
        blocks: &[BytecodeBlock],
    ) {
        let successors = blocks[block.index()].successors().to_vec();
        let mut ready = Vec::new();
        for successor in successors {
            let producers = &mut self.producers.blocks[successor.target.index()];
            if producers.unsealed_predecessors > 0 {
                producers.unsealed_predecessors -= 1;
                if producers.unsealed_predecessors == 0 {
                    ready.push(successor.target);
                }
            }
        }
        for block in ready {
            self.seal_block(block, blocks);
        }
    }

    fn seal_block(&mut self, block: BytecodeBlockId, blocks: &[BytecodeBlock]) {
        if self.producers.blocks[block.index()].sealed {
            return;
        }
        self.producers.blocks[block.index()].sealed = true;
        let incomplete = std::mem::take(&mut self.producers.blocks[block.index()].incomplete_phis);
        for (register, phi) in incomplete {
            self.add_phi_operands(blocks, block, register, phi);
        }
    }

    pub(in crate::graph) fn seal_all_remaining(&mut self, blocks: &[BytecodeBlock]) {
        for index in 0..blocks.len() {
            self.seal_block(BytecodeBlockId::new(index), blocks);
        }
    }

    pub(in crate::graph) fn simplify_phis(&mut self, instructions: &mut [BytecodeInstruction]) {
        let mut replacements = HashMap::new();

        loop {
            let mut changed = false;
            for index in 0..self.phis.len() {
                let id = BytecodePhiId::new(index);
                if replacements.contains_key(&id) {
                    continue;
                }

                let phi = BytecodeOperand::Phi(id);
                let mut value = None;
                let mut nontrivial = false;
                for operand in self.phis[index].operands.iter().copied() {
                    let operand = resolve_operand(operand, &replacements);
                    if operand == phi || value == Some(operand) {
                        continue;
                    }
                    if value.is_some() {
                        nontrivial = true;
                        break;
                    }
                    value = Some(operand);
                }

                if !nontrivial {
                    let replacement =
                        value.unwrap_or_else(|| BytecodeOperand::VmRegister(self.registers[&phi]));
                    replacements.insert(id, replacement);
                    changed = true;
                }
            }

            if !changed {
                break;
            }
        }

        for instruction in instructions {
            for operand in &mut instruction.operands {
                *operand = resolve_operand(*operand, &replacements);
            }
        }
        for phi in &mut self.phis {
            for operand in &mut phi.operands {
                *operand = resolve_operand(*operand, &replacements);
            }
        }
        for phis in &mut self.block_phis {
            phis.retain(|id| !replacements.contains_key(id));
        }
    }

    pub(super) fn producers_up_to_top(
        &mut self,
        blocks: &[BytecodeBlock],
        mut register: Register,
    ) -> Vec<BytecodeOperand> {
        let block = self.current_block;
        let Some(multi_return) = self.producers.blocks[block.index()].multi_return else {
            return Vec::new();
        };
        let multi_return_start = self.producers.blocks[block.index()].multi_return_start;

        let mut operands = Vec::new();
        while register < multi_return_start {
            if let Some(operand) = self.read_variable(blocks, block, register) {
                operands.push(operand);
            }
            register += 1;
        }
        operands.push(multi_return);

        let block_producers = &mut self.producers.blocks[block.index()];
        block_producers.multi_return = None;
        block_producers.multi_return_start = 0xff;
        operands
    }

    pub(super) fn produce(&mut self, register: Register, instruction: BytecodeInstructionId) {
        self.add_producer(register, BytecodeOperand::Instruction(instruction));
    }

    pub(super) fn produce_projection(
        &mut self,
        register: Register,
        instruction: BytecodeInstructionId,
        index: u32,
    ) {
        let projection = self.projection(BytecodeOperand::Instruction(instruction), index);
        self.add_producer(register, projection);
    }

    pub(super) fn projection(&mut self, source: BytecodeOperand, index: u32) -> BytecodeOperand {
        self.projections.push(BytecodeProjection { source, index });
        BytecodeOperand::Projection(BytecodeProjectionId::new(self.projections.len() - 1))
    }

    pub(super) fn add_producer(&mut self, register: Register, operand: BytecodeOperand) {
        self.add_producer_to(self.current_block, register, operand);
    }

    pub(super) fn add_producer_to(
        &mut self,
        block: BytecodeBlockId,
        register: Register,
        operand: BytecodeOperand,
    ) {
        self.registers.insert(operand, register);
        let block_producers = &mut self.producers.blocks[block.index()];
        block_producers.own.insert(register, operand);
        block_producers.invalid_after = block_producers.invalid_after.max(i32::from(register));
    }

    fn is_unreachable(
        &self,
        blocks: &[BytecodeBlock],
        block: BytecodeBlockId,
        visited: &mut HashSet<BytecodeBlockId>,
    ) -> bool {
        if block.index() == 0 {
            return false;
        }
        if !visited.insert(block) {
            return true;
        }
        blocks[block.index()]
            .predecessors()
            .iter()
            .filter(|edge| edge.kind != BytecodeEdgeKind::Loop)
            .all(|edge| self.is_unreachable(blocks, edge.target, visited))
    }

    pub(super) fn apply_call(
        &mut self,
        instruction: BytecodeInstructionId,
        target: Register,
        results: i32,
    ) {
        let block_producers = &mut self.producers.blocks[self.current_block.index()];
        block_producers.own.retain(|register, _| *register < target);
        block_producers
            .cached
            .retain(|register, _| *register < target);

        if results < 0 {
            block_producers.multi_return = Some(BytecodeOperand::Instruction(instruction));
            block_producers.multi_return_start = target;
            block_producers.invalid_after = 255;
            return;
        }

        block_producers.invalid_after = i32::from(target) - 1 + results;
    }
}

fn resolve_operand(
    mut operand: BytecodeOperand,
    replacements: &HashMap<BytecodePhiId, BytecodeOperand>,
) -> BytecodeOperand {
    let mut remaining = replacements.len();
    while let BytecodeOperand::Phi(id) = operand {
        let Some(replacement) = replacements.get(&id).copied() else {
            break;
        };
        operand = replacement;
        remaining -= 1;
        if remaining == 0 {
            break;
        }
    }
    operand
}