package vaporize;

import java.util.*;
import misc.*;

public class RegisterAlloc {

    private LIRDict intervals;
    private String[] all_registers;
    private Stack<String> free_registers;
    private TreeSet<LIRVar> active;
    private int spill_start;

    public RegisterAlloc(LIRDict intervals, String[] all_registers) {
        this.spill_start = 14;
        this.intervals = intervals;
        this.all_registers = all_registers;
        this.free_registers = new Stack<String>();
        this.free_registers.addAll(Arrays.asList(this.all_registers));
        this.active = new TreeSet<LIRVar>((v1, v2) -> {
                int ret;
                if (v1.getLastUse() > v2.getLastUse())
                    ret = 1;
                else if (v1.getLastUse() < v2.getLastUse())
                    ret = -1;
                else if (v1.equals(v2))
                    ret = 0;
                else
                    ret = 1;
                return ret;
        });

        MinimalLogger.info(String.format("Starting allocation with registers %s",
                                         this.free_registers.toString()));

        String register;
        for (LIRVar interval : this.intervals.getIntervals()) {
            this.expireOldIntervals(interval);
            if (this.active.size() >= this.all_registers.length)
                this.spillAtInterval(interval);
            else {
                register = this.free_registers.pop();
                interval.assignRegister(register);
                this.active.add(interval);
                MinimalLogger.info(String.format("Assigning register %s to %s.",
                                                 register,
                                                 interval.toString()));

            }
        }
    }

    private void expireOldIntervals(LIRVar interval) {
        // MinimalLogger.info(String.format("Active: %s",
        //                                  this.active.toString()));
        // MinimalLogger.info(String.format("Available: %s",
        //                                  this.free_registers.toString()));

        Iterator<LIRVar> iterator = this.active.iterator();
        while (iterator.hasNext()) {
            LIRVar active = iterator.next();
            if (active.getLastUse() >= interval.getFirstDef())
                return;
            MinimalLogger.info(String.format("%s register %s expired!",
                                             active.toString(),
                                             active.getAssignedRegister()));
            iterator.remove(); // Remove the element from the original active set
            this.free_registers.push(active.getAssignedRegister());
        }
    }

    private void spillAtInterval(LIRVar interval) {
        this.intervals.addSpilledNum();
        LIRVar spill = this.active.last();
        if (spill.getLastUse() > interval.getLastUse()) {
            interval.assignRegister(spill.getAssignedRegister());
            this.active.add(interval);
            this.active.remove(spill);
            interval = spill;
        }
        MinimalLogger.severe(String.format("Ran out of free registers, had to spill %s!",
                                           interval.toString()));
        interval.assignRegister(String.format("local[%d]", this.spill_start++));
    }

}