<commit_before>package com.miracle.intermediate.visitor.irtranser;

import com.miracle.intermediate.Root;
import com.miracle.intermediate.instruction.Call;
import com.miracle.intermediate.instruction.HeapAllocate;
import com.miracle.intermediate.instruction.Move;
import com.miracle.intermediate.instruction.fork.Return;
import com.miracle.intermediate.number.*;
import com.miracle.intermediate.number.Number;
import com.miracle.intermediate.structure.BasicBlock;
import com.miracle.intermediate.structure.Function;
import com.miracle.intermediate.visitor.BaseIRVisitor;

import java.util.HashSet;
import java.util.Set;
import java.util.Stack;

import static com.miracle.intermediate.number.PhysicalRegister.RAX;

public class LLIRTransformer extends BaseIRVisitor {
    private Set<BasicBlock> blockProcessed;
    private Function curFunction;
    private BasicBlock.Node node;

    private void enroll(Number register) {
        if (register == null) {
            throw new RuntimeException("GG");
        }
        if (register instanceof PhysicalRegister) {
            curFunction.buffer.enroll(((PhysicalRegister) register));
        } else if (register instanceof StackRegister) {
            curFunction.buffer.enroll((StackRegister) register);
        } else if (register instanceof VirtualRegister) {
            enroll(((VirtualRegister) register).getRealName());
        }
    }

    @Override
    public void visit(Root ir) {
        blockProcessed = new HashSet<>();
        ir.globalFunction.forEach((key, value) -> value.accept(this));
    }

    @Override
    public void visit(Function function) {
        curFunction = function;
        function.parameters.forEach(this::enroll);
        curFunction.getEntryBasicBlock().accept(this);
        for (BasicBlock block : curFunction.getPostOrder()) {
            Set<PhysicalRegister> live = new HashSet<PhysicalRegister>() {{
                block.liveliness.liveOut.forEach(element -> {
                    if (element.getRealName() instanceof PhysicalRegister) {
                        add((PhysicalRegister) element.getRealName());
                    }
                });
            }};
            for (node = block.tail.getPrev(); node != block.getHead().getPrev(); node = node.getPrev()) {
                if (node.instruction instanceof Call) {
                    ((Call) node.instruction).callerSave.addAll(
                            ((Call) node.instruction).function.buffer.getCallerSaveRegisters()
                    );
                    ((Call) node.instruction).callerSave.retainAll(live);
                } else if (node.instruction instanceof HeapAllocate) {
                    ((HeapAllocate) node.instruction).callerSave.retainAll(live);
                }
                node.instruction.getDefNumbers().forEach(element -> {
                    if (element instanceof VirtualRegister && ((VirtualRegister) element).getRealName() instanceof PhysicalRegister) {
                        live.remove(((VirtualRegister) element).getRealName());<commit_msg>Register allocator added.<commit_after>package com.miracle.intermediate.visitor.irtranser;

import com.miracle.MiracleOption;
import com.miracle.intermediate.Root;
import com.miracle.intermediate.instruction.Call;
import com.miracle.intermediate.instruction.HeapAllocate;
import com.miracle.intermediate.instruction.Move;
import com.miracle.intermediate.instruction.fork.Return;
import com.miracle.intermediate.number.*;
import com.miracle.intermediate.number.Number;
import com.miracle.intermediate.structure.BasicBlock;
import com.miracle.intermediate.structure.Function;
import com.miracle.intermediate.visitor.BaseIRVisitor;

import java.util.HashSet;
import java.util.Set;
import java.util.Stack;

import static com.miracle.intermediate.number.PhysicalRegister.RAX;

public class LLIRTransformer extends BaseIRVisitor {
    private Set<BasicBlock> blockProcessed;
    private Function curFunction;
    private BasicBlock.Node node;

    private void enroll(Number register) {
        if (register == null) {
            throw new RuntimeException("GG");
        }
        if (register instanceof PhysicalRegister) {
            curFunction.buffer.enroll(((PhysicalRegister) register));
        } else if (register instanceof StackRegister) {
            curFunction.buffer.enroll((StackRegister) register);
        } else if (register instanceof VirtualRegister) {
            enroll(((VirtualRegister) register).getRealName());
        }
    }

    @Override
    public void visit(Root ir) {
        blockProcessed = new HashSet<>();
        ir.globalFunction.forEach((key, value) -> value.accept(this));
    }

    @Override
    public void visit(Function function) {
        curFunction = function;
        function.parameters.forEach(this::enroll);
        curFunction.getEntryBasicBlock().accept(this);
        for (BasicBlock block : curFunction.getPostOrder()) {
            Set<PhysicalRegister> live = new HashSet<PhysicalRegister>() {{
                block.liveliness.liveOut.forEach(element -> {
                    if (element.getRealName() instanceof PhysicalRegister) {
                        add((PhysicalRegister) element.getRealName());
                    }
                });
            }};
            for (node = block.tail.getPrev(); node != block.getHead().getPrev(); node = node.getPrev()) {
                if (node.instruction instanceof Call) {
                    ((Call) node.instruction).callerSave.addAll(
                            ((Call) node.instruction).function.buffer.getCallerSaveRegisters()
                    );
                    for (int i = 0, size = ((Call) node.instruction).parameters.size(); i < size && i < MiracleOption.CallingConvention.size(); i++) {
                        ((Call) node.instruction).callerSave.add(
                                PhysicalRegister.getBy16BITName(MiracleOption.CallingConvention.get(i), ((Call) node.instruction).parameters.get(i).getNumberSize())
                        );
                    }
                    ((Call) node.instruction).callerSave.retainAll(live);
                } else if (node.instruction instanceof HeapAllocate) {
                    ((HeapAllocate) node.instruction).callerSave.retainAll(live);
                }
                node.instruction.getDefNumbers().forEach(element -> {
                    if (element instanceof VirtualRegister && ((VirtualRegister) element).getRealName() instanceof PhysicalRegister) {
                        live.remove(((VirtualRegister) element).getRealName());
{'fraction_non_alphanumeric': 0.10865926038654695, 'fraction_numerical': 0.0004565515142291889, 'mean_word_length': 2.607025246981339, 'pattern_counts': {'":': 0, '<': 13, '<?xml version=': 0, '>': 17, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '2370514', 'n_tokens_mistral': 1702, 'n_tokens_neox': 1619, 'n_words': 310}