register_alloc.sc

#   register allocation algorithm idea in O(n)

    by leonard ritter

    based on DAG visualization here:

    https://twitter.com/paniq/status/1297661746832973825

using import Array
using import Map

# if false, find lifetime from out edges instead - either method has
    its advantages.
let USE_IN_EDGES = true

# the hardcoded demo DAG is based on the petersen graph

fn in-edges (i)
    inline node (...)
        local result : (Array i32)
        va-map
            inline (x)
                'append result x
            ...
        result
    switch i
    case 0
        node 1 4
    case 1
        node 2
    case 2
        node 3
    case 3
        node;
    case 4
        node 3
    case 5
        node 0
    case 6
        node 1 8
    case 7
        node 2 5
    case 8
        node 3 5
    case 9
        node 4 6 7
    default
        return ((Array i32))

fn out-edges (i)
    inline nodeto (...)
        local result : (Array i32)
        va-map
            inline (x)
                'append result x
            ...
        result
    switch i
    case 0
        nodeto 5
    case 1
        nodeto 0 6
    case 2
        nodeto 1 7
    case 3
        nodeto 2 4 8
    case 4
        nodeto 0 9
    case 5
        nodeto 7 8
    case 6
        nodeto 9
    case 7
        nodeto 9
    case 8
        nodeto 6
    case 9
        nodeto;
    default
        return ((Array i32))

# hardcoded topo order as topological sorting is not what we want to demo here
local topo-order =
    arrayof i32 3 2 4 1 0 5 7 8 6 9
local inverse-topo-order : (array i32 10)
local node-eol : (array i32 10)
# build map inverse
for i node in (enumerate topo-order)
    static-if USE_IN_EDGES
        # keep incoming nodes alive for as long as we're alive
        for source in (in-edges node)
            let dest = (node-eol @ source)
            dest = (max dest i)
    else
        inverse-topo-order @ node = i

# now do register allocation
let MAX_REGISTERS = 10
local registers : (array i32 MAX_REGISTERS)
# walk nodes in topological order
for i node in (enumerate topo-order)
    let register =
        # find free register
        for j end-of-life in (enumerate registers)
            # any register being consumed in this frame or earlier is fine
            if (end-of-life <= i)
                # found one
                break j
        else
            error "register overflow"
    let end-of-life =
        static-if USE_IN_EDGES
            node-eol @ node
        else
            # find last access
            fold (best = 0) for j in (out-edges node)
                let time = (inverse-topo-order @ j)
                max best time
    # allocate register providing end of life time
    registers @ register = end-of-life
    # the final assigned register
    print "node" node "sets register" register