Research GitLab has introduced a user quota limitation. The new rule limits each user to have 50 Gb. The quota doesn't restrict group projects. If you have any concern with this, please talk to CECS Gitlab Admin at N110 (b) CSIT building.

coloring.rs 26.4 KB
Newer Older
Isaac Oscar Gariano's avatar
Isaac Oscar Gariano committed
1
// Copyright 2017 The Australian National University
qinsoon's avatar
qinsoon committed
2
//
Isaac Oscar Gariano's avatar
Isaac Oscar Gariano committed
3 4 5
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
qinsoon's avatar
qinsoon committed
6
//
Isaac Oscar Gariano's avatar
Isaac Oscar Gariano committed
7
//     http://www.apache.org/licenses/LICENSE-2.0
qinsoon's avatar
qinsoon committed
8
//
Isaac Oscar Gariano's avatar
Isaac Oscar Gariano committed
9 10 11 12 13 14
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

15 16
extern crate hprof;

qinsoon's avatar
qinsoon committed
17
use ast::ptr::*;
qinsoon's avatar
qinsoon committed
18 19 20
use ast::ir::*;
use compiler::backend;
use compiler::backend::reg_alloc::graph_coloring;
21
use compiler::backend::reg_alloc::graph_coloring::liveness::InterferenceGraph;
qinsoon's avatar
qinsoon committed
22 23
use compiler::machine_code::CompiledFunction;
use vm::VM;
24

25
use utils::vec_utils;
26
use utils::LinkedHashSet;
27
use utils::LinkedHashMap;
qinsoon's avatar
qinsoon committed
28

qinsoon's avatar
qinsoon committed
29
use std::cell::RefCell;
qinsoon's avatar
qinsoon committed
30

31 32 33
use compiler::backend::reg_alloc::graph_coloring::liveness::Move;
use compiler::backend::reg_alloc::graph_coloring::petgraph::graph::NodeIndex;

qinsoon's avatar
qinsoon committed
34 35
const COALESCING: bool = true;
const MAX_REWRITE_ITERATIONS_ALLOWED: usize = 10;
36

qinsoon's avatar
qinsoon committed
37 38
/// GraphColoring algorithm
/// based on Appel's book section 11.4
qinsoon's avatar
qinsoon committed
39
pub struct GraphColoring<'a> {
qinsoon's avatar
qinsoon committed
40
    // context
qinsoon's avatar
qinsoon committed
41 42 43
    pub func: &'a mut MuFunctionVersion,
    pub cf: &'a mut CompiledFunction,
    pub vm: &'a VM,
44
    pub ig: InterferenceGraph,
qinsoon's avatar
qinsoon committed
45

46 47 48 49
    /// how many coloring iteration have we done?
    /// In case that a bug may trigger the coloring iterate endlessly, we use this count to stop
    iteration_count: usize,

qinsoon's avatar
qinsoon committed
50
    /// machine registers, preassigned a color
51
    precolored: LinkedHashSet<MuID>,
qinsoon's avatar
qinsoon committed
52
    /// all colors available
53
    colors: LinkedHashMap<backend::RegGroup, LinkedHashSet<MuID>>,
qinsoon's avatar
qinsoon committed
54
    /// temporaries, not precolored and not yet processed
55
    initial: Vec<MuID>,
qinsoon's avatar
qinsoon committed
56 57

    /// list of low-degree non-move-related nodes
58
    worklist_simplify: LinkedHashSet<MuID>,
qinsoon's avatar
qinsoon committed
59
    /// low-degree move related nodes
60
    worklist_freeze: LinkedHashSet<MuID>,
qinsoon's avatar
qinsoon committed
61
    /// nodes marked for spilling during this round
62
    worklist_spill: Vec<MuID>,
qinsoon's avatar
qinsoon committed
63
    /// nodes marked for spilling during this round
64
    spilled_nodes: Vec<MuID>,
qinsoon's avatar
qinsoon committed
65 66
    /// temps that have been coalesced
    /// when u <- v is coalesced, v is added to this set and u put back on some work list
67
    coalesced_nodes: LinkedHashSet<MuID>,
qinsoon's avatar
qinsoon committed
68
    /// nodes successfully colored
69
    colored_nodes: Vec<MuID>,
qinsoon's avatar
qinsoon committed
70
    /// stack containing temporaries removed from the graph
71
    select_stack: Vec<MuID>,
qinsoon's avatar
qinsoon committed
72 73

    /// moves that have been coalesced
74
    coalesced_moves: LinkedHashSet<Move>,
qinsoon's avatar
qinsoon committed
75
    /// moves whose source and target interfere
76
    constrained_moves: LinkedHashSet<Move>,
qinsoon's avatar
qinsoon committed
77 78 79 80 81 82 83 84
    /// moves that will no longer be considered for coalescing
    frozen_moves: LinkedHashSet<Move>,
    /// moves enabled for possible coalescing
    worklist_moves: Vec<Move>,
    /// moves not yet ready for coalescing
    active_moves: LinkedHashSet<Move>,

    /// a mapping from a node to the list of moves it is associated with
85
    movelist: LinkedHashMap<MuID, LinkedHashSet<Move>>,
qinsoon's avatar
qinsoon committed
86
    /// when a move (u, v) has been coalesced, and v put in coalescedNodes, then alias(v) = u
87
    alias: LinkedHashMap<MuID, MuID>,
qinsoon's avatar
qinsoon committed
88

qinsoon's avatar
qinsoon committed
89 90 91 92
    // for validation use
    /// we need to log all registers get spilled with their spill location
    spill_history: LinkedHashMap<MuID, P<Value>>,
    /// we need to know the mapping between scratch temp -> original temp
93
    spill_scratch_temps: LinkedHashMap<MuID, MuID>
qinsoon's avatar
qinsoon committed
94 95
}

qinsoon's avatar
qinsoon committed
96
impl<'a> GraphColoring<'a> {
qinsoon's avatar
qinsoon committed
97
    /// starts coloring
qinsoon's avatar
qinsoon committed
98 99 100
    pub fn start(
        func: &'a mut MuFunctionVersion,
        cf: &'a mut CompiledFunction,
101
        vm: &'a VM
qinsoon's avatar
qinsoon committed
102 103 104 105 106 107 108
    ) -> GraphColoring<'a> {
        GraphColoring::start_with_spill_history(
            LinkedHashMap::new(),
            LinkedHashMap::new(),
            0,
            func,
            cf,
109
            vm
qinsoon's avatar
qinsoon committed
110
        )
qinsoon's avatar
qinsoon committed
111 112
    }

qinsoon's avatar
qinsoon committed
113
    /// restarts coloring with spill history
qinsoon's avatar
qinsoon committed
114 115 116 117 118 119
    fn start_with_spill_history(
        spill_history: LinkedHashMap<MuID, P<Value>>,
        spill_scratch_temps: LinkedHashMap<MuID, MuID>,
        iteration_count: usize,
        func: &'a mut MuFunctionVersion,
        cf: &'a mut CompiledFunction,
120
        vm: &'a VM
qinsoon's avatar
qinsoon committed
121 122 123
    ) -> GraphColoring<'a> {
        assert!(
            iteration_count < MAX_REWRITE_ITERATIONS_ALLOWED,
124
            "reach graph coloring max rewrite iterations ({}), probably something is going wrong",
qinsoon's avatar
qinsoon committed
125 126
            MAX_REWRITE_ITERATIONS_ALLOWED
        );
127 128
        let iteration_count = iteration_count + 1;

129
        trace!("Initializing coloring allocator...");
qinsoon's avatar
qinsoon committed
130 131
        cf.mc().trace_mc();

qinsoon's avatar
qinsoon committed
132 133
        let ig = graph_coloring::build_inteference_graph(cf, func);

134
        let coloring = GraphColoring {
qinsoon's avatar
qinsoon committed
135 136 137 138
            func: func,
            cf: cf,
            vm: vm,
            ig: ig,
139
            iteration_count: iteration_count,
140
            precolored: LinkedHashSet::new(),
141
            colors: {
142
                let mut map = LinkedHashMap::new();
143 144
                map.insert(backend::RegGroup::GPR, LinkedHashSet::new());
                map.insert(backend::RegGroup::FPR, LinkedHashSet::new());
145 146 147
                map
            },
            colored_nodes: Vec::new(),
148 149
            initial: Vec::new(),
            worklist_moves: Vec::new(),
150
            movelist: LinkedHashMap::new(),
151 152 153 154
            active_moves: LinkedHashSet::new(),
            coalesced_nodes: LinkedHashSet::new(),
            coalesced_moves: LinkedHashSet::new(),
            constrained_moves: LinkedHashSet::new(),
155
            alias: LinkedHashMap::new(),
156
            worklist_spill: Vec::new(),
157
            spilled_nodes: Vec::new(),
qinsoon's avatar
qinsoon committed
158
            spill_history: spill_history,
qinsoon's avatar
qinsoon committed
159
            spill_scratch_temps: spill_scratch_temps,
160 161 162
            worklist_freeze: LinkedHashSet::new(),
            frozen_moves: LinkedHashSet::new(),
            worklist_simplify: LinkedHashSet::new(),
163
            select_stack: Vec::new()
qinsoon's avatar
qinsoon committed
164
        };
qinsoon's avatar
qinsoon committed
165

qinsoon's avatar
qinsoon committed
166
        coloring.regalloc()
qinsoon's avatar
qinsoon committed
167
    }
qinsoon's avatar
qinsoon committed
168

qinsoon's avatar
qinsoon committed
169
    /// returns formatted string for an ID
qinsoon's avatar
qinsoon committed
170 171 172 173
    fn display_id(&self, id: MuID) -> String {
        self.func.context.get_temp_display(id)
    }

qinsoon's avatar
qinsoon committed
174
    /// returns formatted string for a move
qinsoon's avatar
qinsoon committed
175
    fn display_move(&self, m: Move) -> String {
qinsoon's avatar
qinsoon committed
176 177
        format!(
            "Move: {} -> {}",
178 179
            self.display_id(m.from),
            self.display_id(m.to)
qinsoon's avatar
qinsoon committed
180
        )
qinsoon's avatar
qinsoon committed
181
    }
qinsoon's avatar
qinsoon committed
182 183

    /// does coloring register allocation
184
    fn regalloc(mut self) -> GraphColoring<'a> {
qinsoon's avatar
qinsoon committed
185
        trace!("---InterenceGraph---");
qinsoon's avatar
qinsoon committed
186
        self.ig.print(&self.func.context);
qinsoon's avatar
qinsoon committed
187 188 189

        // start timing for graph coloring
        let _p = hprof::enter("regalloc: graph coloring");
qinsoon's avatar
qinsoon committed
190

qinsoon's avatar
qinsoon committed
191
        // precolor for all machine registers
192
        for reg in backend::all_regs().values() {
193
            let reg_id = reg.extract_ssa_id().unwrap();
194
            self.precolored.insert(reg_id);
195
        }
qinsoon's avatar
qinsoon committed
196

qinsoon's avatar
qinsoon committed
197
        // put usable registers as available colors
198
        for reg in backend::all_usable_regs().iter() {
qinsoon's avatar
qinsoon committed
199
            let reg_id = reg.extract_ssa_id().unwrap();
qinsoon's avatar
qinsoon committed
200 201
            let group = backend::pick_group_for_reg(reg_id);
            self.colors.get_mut(&group).unwrap().insert(reg_id);
qinsoon's avatar
qinsoon committed
202
        }
qinsoon's avatar
qinsoon committed
203 204

        // push uncolored nodes to initial work set
qinsoon's avatar
qinsoon committed
205 206
        for node in self.ig.nodes() {
            if !self.ig.is_colored(node) {
207
                self.initial.push(node);
qinsoon's avatar
qinsoon committed
208 209
            }
        }
qinsoon's avatar
qinsoon committed
210 211

        // initialize work
qinsoon's avatar
qinsoon committed
212 213
        self.build();
        self.make_work_list();
qinsoon's avatar
qinsoon committed
214 215

        // main loop
qinsoon's avatar
qinsoon committed
216 217 218 219 220 221 222 223 224 225
        while {
            if !self.worklist_simplify.is_empty() {
                self.simplify();
            } else if !self.worklist_moves.is_empty() {
                self.coalesce();
            } else if !self.worklist_freeze.is_empty() {
                self.freeze();
            } else if !self.worklist_spill.is_empty() {
                self.select_spill();
            }
qinsoon's avatar
qinsoon committed
226 227 228

            !(self.worklist_simplify.is_empty() && self.worklist_moves.is_empty() &&
                  self.worklist_freeze.is_empty() && self.worklist_spill.is_empty())
qinsoon's avatar
qinsoon committed
229
        } {}
qinsoon's avatar
qinsoon committed
230 231

        // pick color for nodes
qinsoon's avatar
qinsoon committed
232 233
        self.assign_colors();

qinsoon's avatar
qinsoon committed
234
        // finish
235 236
        drop(_p);

qinsoon's avatar
qinsoon committed
237
        // if we need to spill
qinsoon's avatar
qinsoon committed
238 239
        if !self.spilled_nodes.is_empty() {
            trace!("spill required");
qinsoon's avatar
qinsoon committed
240 241 242
            if cfg!(debug_assertions) {
                trace!("nodes to be spilled:");
                for node in self.spilled_nodes.iter() {
243
                    trace!("{}", *node);
qinsoon's avatar
qinsoon committed
244 245 246
                }
            }

qinsoon's avatar
qinsoon committed
247
            // rewrite program to insert spilling code
qinsoon's avatar
qinsoon committed
248
            self.rewrite_program();
qinsoon's avatar
qinsoon committed
249

qinsoon's avatar
qinsoon committed
250
            // recursively redo graph coloring
qinsoon's avatar
qinsoon committed
251 252 253 254 255 256
            return GraphColoring::start_with_spill_history(
                self.spill_history.clone(),
                self.spill_scratch_temps.clone(),
                self.iteration_count,
                self.func,
                self.cf,
257
                self.vm
qinsoon's avatar
qinsoon committed
258
            );
qinsoon's avatar
qinsoon committed
259 260
        }

261
        self
qinsoon's avatar
qinsoon committed
262
    }
qinsoon's avatar
qinsoon committed
263

264 265 266 267 268 269 270 271 272 273 274 275 276 277 278 279 280 281 282 283 284 285 286
    fn add_to_movelist(
        movelist: &mut LinkedHashMap<MuID, LinkedHashSet<Move>>,
        reg: MuID,
        mov: Move
    ) {
        if movelist.contains_key(&reg) {
            let mut list = movelist.get_mut(&reg).unwrap();
            list.insert(mov);
        } else {
            let mut list = LinkedHashSet::new();
            list.insert(mov);
            movelist.insert(reg, list);
        }
    }

    fn get_movelist(&self, reg: MuID) -> LinkedHashSet<Move> {
        if let Some(list) = self.movelist.get(&reg) {
            list.clone()
        } else {
            LinkedHashSet::new()
        }
    }

qinsoon's avatar
qinsoon committed
287
    fn build(&mut self) {
qinsoon's avatar
qinsoon committed
288 289
        if COALESCING {
            trace!("coalescing enabled, build move list");
qinsoon's avatar
qinsoon committed
290
            let ref ig = self.ig;
291
            for m in ig.moves().iter() {
qinsoon's avatar
qinsoon committed
292
                trace!("add to movelist: {:?}", m);
293 294 295 296
                self.worklist_moves.push(*m);

                GraphColoring::add_to_movelist(&mut self.movelist, m.from, *m);
                GraphColoring::add_to_movelist(&mut self.movelist, m.to, *m);
qinsoon's avatar
qinsoon committed
297
            }
qinsoon's avatar
qinsoon committed
298 299
        } else {
            trace!("coalescing disabled");
qinsoon's avatar
qinsoon committed
300 301
        }
    }
qinsoon's avatar
qinsoon committed
302

qinsoon's avatar
qinsoon committed
303
    fn make_work_list(&mut self) {
304
        trace!("Making work list from initials...");
qinsoon's avatar
qinsoon committed
305
        while !self.initial.is_empty() {
306
            let node = self.initial.pop().unwrap();
qinsoon's avatar
qinsoon committed
307

qinsoon's avatar
qinsoon committed
308 309
            if {
                // condition: degree >= K
qinsoon's avatar
qinsoon committed
310
                let degree = self.ig.get_degree_of(node);
311
                let n_regs = self.n_regs_for_node(node);
qinsoon's avatar
qinsoon committed
312

qinsoon's avatar
qinsoon committed
313 314
                degree >= n_regs
            } {
qinsoon's avatar
qinsoon committed
315 316
                trace!(
                    "{} 's degree >= reg number limit (K), push to spill list",
317
                    node
qinsoon's avatar
qinsoon committed
318
                );
319
                self.worklist_spill.push(node);
qinsoon's avatar
qinsoon committed
320
            } else if self.is_move_related(node) {
321
                trace!("{} is move related, push to freeze list", node);
qinsoon's avatar
qinsoon committed
322 323
                self.worklist_freeze.insert(node);
            } else {
qinsoon's avatar
qinsoon committed
324 325
                trace!(
                    "{} has small degree and not move related, push to simplify list",
326
                    node
qinsoon's avatar
qinsoon committed
327
                );
qinsoon's avatar
qinsoon committed
328 329 330 331
                self.worklist_simplify.insert(node);
            }
        }
    }
qinsoon's avatar
qinsoon committed
332

333
    fn n_regs_for_node(&self, node: MuID) -> usize {
334
        backend::number_of_usable_regs_in_group(self.ig.get_group_of(node))
335
    }
qinsoon's avatar
qinsoon committed
336

337
    fn is_move_related(&mut self, node: MuID) -> bool {
qinsoon's avatar
qinsoon committed
338 339
        !self.node_moves(node).is_empty()
    }
340 341 342 343 344 345 346 347 348 349

    fn is_spillable(&self, temp: MuID) -> bool {
        // if a temporary is created as scratch temp for a spilled temporary, we
        // should not spill it again (infinite loop otherwise)
        if self.spill_scratch_temps.contains_key(&temp) {
            false
        } else {
            true
        }
    }
qinsoon's avatar
qinsoon committed
350

351
    fn node_moves(&mut self, node: MuID) -> LinkedHashSet<Move> {
352
        let mut moves = LinkedHashSet::new();
qinsoon's avatar
qinsoon committed
353

qinsoon's avatar
qinsoon committed
354 355 356 357
        // addAll(active_moves)
        for m in self.active_moves.iter() {
            moves.insert(m.clone());
        }
qinsoon's avatar
qinsoon committed
358

qinsoon's avatar
qinsoon committed
359 360 361 362
        // addAll(worklist_moves)
        for m in self.worklist_moves.iter() {
            moves.insert(m.clone());
        }
qinsoon's avatar
qinsoon committed
363

364
        let mut retained = LinkedHashSet::new();
365
        let movelist = self.get_movelist(node);
qinsoon's avatar
qinsoon committed
366
        for m in moves.iter() {
367
            if movelist.contains(m) {
qinsoon's avatar
qinsoon committed
368 369 370
                retained.insert(*m);
            }
        }
qinsoon's avatar
qinsoon committed
371

qinsoon's avatar
qinsoon committed
372 373
        retained
    }
qinsoon's avatar
qinsoon committed
374

375
    fn simplify(&mut self) {
376
        // remove next element from worklist_simplify, we know its not empty
377
        let node = self.worklist_simplify.pop_front().unwrap();
qinsoon's avatar
qinsoon committed
378

379
        trace!("Simplifying {}", node);
qinsoon's avatar
qinsoon committed
380

381
        self.select_stack.push(node);
382
        for m in self.adjacent(node).iter() {
383
            self.decrement_degree(*m);
384 385
        }
    }
qinsoon's avatar
qinsoon committed
386

387
    fn adjacent(&self, n: MuID) -> LinkedHashSet<MuID> {
388
        let mut adj = LinkedHashSet::new();
qinsoon's avatar
qinsoon committed
389

390
        // add n's successors
391 392
        for s in self.ig.get_adj_list(n).iter() {
            adj.insert(*s);
393
        }
qinsoon's avatar
qinsoon committed
394

395 396 397 398
        // removeAll(select_stack)
        for s in self.select_stack.iter() {
            adj.remove(s);
        }
qinsoon's avatar
qinsoon committed
399

400 401 402 403
        // removeAll(coalesced_nodes)
        for s in self.coalesced_nodes.iter() {
            adj.remove(s);
        }
qinsoon's avatar
qinsoon committed
404

405 406
        adj
    }
qinsoon's avatar
qinsoon committed
407

408
    fn decrement_degree(&mut self, n: MuID) {
409 410 411
        if self.precolored.contains(&n) {
            return;
        }
qinsoon's avatar
qinsoon committed
412

413
        trace!("decrement degree of {}", n);
qinsoon's avatar
qinsoon committed
414

415
        let d = self.ig.get_degree_of(n);
416
        debug_assert!(d != 0);
417
        self.ig.set_degree_of(n, d - 1);
qinsoon's avatar
qinsoon committed
418

419
        if d == self.n_regs_for_node(n) {
420
            trace!("{}'s degree is K, no longer need to spill it", n);
421 422
            let mut nodes = self.adjacent(n);
            nodes.insert(n);
qinsoon's avatar
qinsoon committed
423
            trace!("enable moves of {:?}", nodes);
424
            self.enable_moves(nodes);
qinsoon's avatar
qinsoon committed
425

426
            vec_utils::remove_value(&mut self.worklist_spill, n);
qinsoon's avatar
qinsoon committed
427

428
            if self.is_move_related(n) {
429
                trace!("{} is move related, push to freeze list", n);
430 431
                self.worklist_freeze.insert(n);
            } else {
432
                trace!("{} is not move related, push to simplify list", n);
433 434 435 436
                self.worklist_simplify.insert(n);
            }
        }
    }
qinsoon's avatar
qinsoon committed
437

438
    fn enable_moves(&mut self, nodes: LinkedHashSet<MuID>) {
439 440 441 442
        for n in nodes.iter() {
            let n = *n;
            for mov in self.node_moves(n).iter() {
                let mov = *mov;
443
                if self.active_moves.contains(&mov) {
qinsoon's avatar
qinsoon committed
444
                    self.active_moves.remove(&mov);
445 446 447 448 449
                    self.worklist_moves.push(mov);
                }
            }
        }
    }
qinsoon's avatar
qinsoon committed
450

451
    fn coalesce(&mut self) {
qinsoon's avatar
qinsoon committed
452
        let m = self.worklist_moves.pop().unwrap();
qinsoon's avatar
qinsoon committed
453

qinsoon's avatar
qinsoon committed
454
        trace!("Coalescing on {}", self.display_move(m));
qinsoon's avatar
qinsoon committed
455 456 457

        // if they are not from the same register group, we cannot coalesce them
        if self.ig.get_group_of(m.from) != self.ig.get_group_of(m.to) {
458
            info!("a move instruction of two temporaries of different register groups");
qinsoon's avatar
qinsoon committed
459 460
            info!("from: {:?}, to: {:?}", m.from, m.to);

qinsoon's avatar
qinsoon committed
461 462
            return;
        }
qinsoon's avatar
qinsoon committed
463

qinsoon's avatar
qinsoon committed
464 465
        let x = self.get_alias(m.from);
        let y = self.get_alias(m.to);
466
        trace!("resolve alias: from {} to {}", x, y);
qinsoon's avatar
qinsoon committed
467

qinsoon's avatar
qinsoon committed
468 469 470 471 472 473
        let (u, v, precolored_u, precolored_v) = {
            if self.precolored.contains(&y) {
                let u = y;
                let v = x;
                let precolored_u = true;
                let precolored_v = self.precolored.contains(&v);
qinsoon's avatar
qinsoon committed
474

qinsoon's avatar
qinsoon committed
475 476 477 478 479 480
                (u, v, precolored_u, precolored_v)
            } else {
                let u = x;
                let v = y;
                let precolored_u = self.precolored.contains(&u);
                let precolored_v = self.precolored.contains(&v);
qinsoon's avatar
qinsoon committed
481

qinsoon's avatar
qinsoon committed
482 483 484
                (u, v, precolored_u, precolored_v)
            }
        };
qinsoon's avatar
qinsoon committed
485 486
        trace!(
            "u={}, v={}, precolored_u={}, precolroed_v={}",
487 488
            u,
            v,
qinsoon's avatar
qinsoon committed
489 490 491 492
            precolored_u,
            precolored_v
        );

493 494 495 496 497 498 499 500 501 502 503 504 505 506 507 508
        // if u or v is a machine register that is not usable/not a color, we cannot coalesce
        if self.precolored.contains(&u) {
            let reg_u = self.ig.get_temp_of(u);
            let group = backend::pick_group_for_reg(reg_u);
            if !self.colors.get(&group).unwrap().contains(&reg_u) {
                return;
            }
        }
        if self.precolored.contains(&v) {
            let reg_v = self.ig.get_temp_of(v);
            let group = backend::pick_group_for_reg(reg_v);
            if !self.colors.get(&group).unwrap().contains(&reg_v) {
                return;
            }
        }

qinsoon's avatar
qinsoon committed
509
        if u == v {
qinsoon's avatar
qinsoon committed
510
            trace!("u == v, coalesce the move");
qinsoon's avatar
qinsoon committed
511 512 513 514
            self.coalesced_moves.insert(m);
            if !precolored_u {
                self.add_worklist(u);
            }
515
        } else if precolored_v || self.ig.is_in_adj_set(u, v) {
516
            trace!("precolored_v: {}", precolored_v);
517
            trace!("is_adj(u, v): {}", self.ig.is_in_adj_set(u, v));
qinsoon's avatar
qinsoon committed
518 519
            trace!("v is precolored or u,v is adjacent, the move is constrained");
            self.constrained_moves.insert(m);
qinsoon's avatar
qinsoon committed
520 521 522 523 524 525
            if !precolored_u {
                self.add_worklist(u);
            }
            if !precolored_v {
                self.add_worklist(v);
            }
qinsoon's avatar
qinsoon committed
526
        } else if (precolored_u && self.ok(u, v)) || (!precolored_u && self.conservative(u, v)) {
qinsoon's avatar
qinsoon committed
527 528 529
            trace!("ok(u, v) = {}", self.ok(u, v));
            trace!("conservative(u, v) = {}", self.conservative(u, v));

qinsoon's avatar
qinsoon committed
530 531 532 533
            trace!(
                "precolored_u&&ok(u,v) || !precolored_u&&conserv(u,v), \
                 coalesce and combine the move"
            );
qinsoon's avatar
qinsoon committed
534 535 536 537 538 539
            self.coalesced_moves.insert(m);
            self.combine(u, v);
            if !precolored_u {
                self.add_worklist(u);
            }
        } else {
qinsoon's avatar
qinsoon committed
540
            trace!("cannot coalesce the move");
qinsoon's avatar
qinsoon committed
541 542 543
            self.active_moves.insert(m);
        }
    }
qinsoon's avatar
qinsoon committed
544

545
    pub fn get_alias(&self, node: MuID) -> MuID {
qinsoon's avatar
qinsoon committed
546 547 548 549 550 551
        if self.coalesced_nodes.contains(&node) {
            self.get_alias(*self.alias.get(&node).unwrap())
        } else {
            node
        }
    }
qinsoon's avatar
qinsoon committed
552

553 554
    fn add_worklist(&mut self, node: MuID) {
        if !self.is_move_related(node) && self.ig.get_degree_of(node) < self.n_regs_for_node(node) {
qinsoon's avatar
qinsoon committed
555 556 557 558
            self.worklist_freeze.remove(&node);
            self.worklist_simplify.insert(node);
        }
    }
qinsoon's avatar
qinsoon committed
559

560
    fn ok(&self, u: MuID, v: MuID) -> bool {
561 562
        for t in self.adjacent(v).iter() {
            let t = *t;
563 564
            if !(self.ig.get_degree_of(t) < self.n_regs_for_node(t) ||
                     self.precolored.contains(&t) || self.ig.is_in_adj_set(t, u))
qinsoon's avatar
qinsoon committed
565
            {
qinsoon's avatar
qinsoon committed
566
                return false;
567
            }
qinsoon's avatar
qinsoon committed
568
        }
qinsoon's avatar
qinsoon committed
569

qinsoon's avatar
qinsoon committed
570 571
        true
    }
qinsoon's avatar
qinsoon committed
572

573
    fn conservative(&self, u: MuID, v: MuID) -> bool {
qinsoon's avatar
qinsoon committed
574 575
        let adj_u = self.adjacent(u);
        let adj_v = self.adjacent(v);
576 577 578 579 580
        let nodes = {
            let mut ret = adj_u;
            ret.add_all(adj_v);
            ret
        };
qinsoon's avatar
qinsoon committed
581

qinsoon's avatar
qinsoon committed
582 583
        let n_reg_for_group = self.n_regs_for_node(u);

qinsoon's avatar
qinsoon committed
584
        let mut k = 0;
585
        for n in nodes.iter() {
586
            if self.precolored.contains(n) || self.ig.get_degree_of(*n) >= n_reg_for_group {
qinsoon's avatar
qinsoon committed
587 588 589
                k += 1;
            }
        }
qinsoon's avatar
qinsoon committed
590

qinsoon's avatar
qinsoon committed
591
        k < n_reg_for_group
qinsoon's avatar
qinsoon committed
592
    }
qinsoon's avatar
qinsoon committed
593

594
    fn combine(&mut self, u: MuID, v: MuID) {
qinsoon's avatar
qinsoon committed
595 596 597 598 599 600 601
        if self.worklist_freeze.contains(&v) {
            self.worklist_freeze.remove(&v);
            self.coalesced_nodes.insert(v);
        } else {
            vec_utils::remove_value(&mut self.worklist_spill, v);
            self.coalesced_nodes.insert(v);
        }
qinsoon's avatar
qinsoon committed
602 603 604

        self.alias.insert(v, u);

qinsoon's avatar
qinsoon committed
605
        {
606 607
            // movelist[u] <- movelist[u] + movelist[v]
            let movelist_v = self.get_movelist(v);
qinsoon's avatar
qinsoon committed
608

609 610 611
            for m in movelist_v.iter() {
                GraphColoring::add_to_movelist(&mut self.movelist, u, *m)
            }
qinsoon's avatar
qinsoon committed
612
        }
qinsoon's avatar
qinsoon committed
613

614
        let mut nodes = LinkedHashSet::new();
qinsoon's avatar
qinsoon committed
615 616
        nodes.insert(v);
        self.enable_moves(nodes);
qinsoon's avatar
qinsoon committed
617

618 619
        for t in self.adjacent(v).iter() {
            let t = *t;
qinsoon's avatar
qinsoon committed
620 621 622
            self.add_edge(t, u);
            self.decrement_degree(t);
        }
qinsoon's avatar
qinsoon committed
623

624 625 626
        if self.worklist_freeze.contains(&u) &&
            self.ig.get_degree_of(u) >= self.n_regs_for_node(u)
        {
qinsoon's avatar
qinsoon committed
627 628 629 630
            self.worklist_freeze.remove(&u);
            self.worklist_spill.push(u);
        }
    }
qinsoon's avatar
qinsoon committed
631

632 633
    fn add_edge(&mut self, u: MuID, v: MuID) {
        if u != v && !self.ig.is_in_adj_set(u, v) {
qinsoon's avatar
qinsoon committed
634
            if !self.precolored.contains(&u) {
635 636 637
                self.ig.add_edge(u, v);
                let degree_u = self.ig.get_degree_of(u);
                self.ig.set_degree_of(u, degree_u + 1);
qinsoon's avatar
qinsoon committed
638 639
            }
            if !self.precolored.contains(&v) {
640 641 642
                self.ig.add_edge(v, u);
                let degree_v = self.ig.get_degree_of(v);
                self.ig.set_degree_of(v, degree_v + 1);
qinsoon's avatar
qinsoon committed
643 644
            }
        }
645
    }
qinsoon's avatar
qinsoon committed
646

647
    fn freeze(&mut self) {
qinsoon's avatar
qinsoon committed
648
        // it is not empty (checked before)
649
        let node = self.worklist_freeze.pop_front().unwrap();
650
        trace!("Freezing {}...", node);
qinsoon's avatar
qinsoon committed
651

qinsoon's avatar
qinsoon committed
652 653 654
        self.worklist_simplify.insert(node);
        self.freeze_moves(node);
    }
qinsoon's avatar
qinsoon committed
655

656
    fn freeze_moves(&mut self, u: MuID) {
657 658
        for m in self.node_moves(u).iter() {
            let m = *m;
qinsoon's avatar
qinsoon committed
659 660 661 662
            let mut v = self.get_alias(m.from);
            if v == self.get_alias(u) {
                v = self.get_alias(m.to);
            }
qinsoon's avatar
qinsoon committed
663

qinsoon's avatar
qinsoon committed
664 665
            self.active_moves.remove(&m);
            self.frozen_moves.insert(m);
qinsoon's avatar
qinsoon committed
666 667

            if !self.precolored.contains(&v) && self.node_moves(v).is_empty() &&
668
                self.ig.get_degree_of(v) < self.n_regs_for_node(v)
qinsoon's avatar
qinsoon committed
669
            {
qinsoon's avatar
qinsoon committed
670 671 672 673
                self.worklist_freeze.remove(&v);
                self.worklist_simplify.insert(v);
            }
        }
674
    }
qinsoon's avatar
qinsoon committed
675

676
    fn select_spill(&mut self) {
qinsoon's avatar
qinsoon committed
677
        trace!("Selecting a node to spill...");
678
        let mut m: Option<MuID> = None;
qinsoon's avatar
qinsoon committed
679

680 681 682 683 684
        for n in self.worklist_spill.iter() {
            let n = *n;
            if m.is_none() {
                m = Some(n);
            } else if {
qinsoon's avatar
qinsoon committed
685 686 687 688
                       // m is not none
                       let temp = self.ig.get_temp_of(m.unwrap());
                       !self.is_spillable(temp)
                   } {
689
                m = Some(n);
690 691 692
            } else if (self.ig.get_spill_cost(n) / (self.ig.get_degree_of(n) as f32)) <
                       (self.ig.get_spill_cost(m.unwrap()) /
                            (self.ig.get_degree_of(m.unwrap()) as f32))
qinsoon's avatar
qinsoon committed
693
            {
694 695 696
                m = Some(n);
            }
        }
qinsoon's avatar
qinsoon committed
697

698 699
        // m is not none
        let m = m.unwrap();
700
        trace!("Spilling {}...", m);
qinsoon's avatar
qinsoon committed
701

702 703 704
        vec_utils::remove_value(&mut self.worklist_spill, m);
        self.worklist_simplify.insert(m);
        self.freeze_moves(m);
705
    }
qinsoon's avatar
qinsoon committed
706

707
    fn assign_colors(&mut self) {
qinsoon's avatar
qinsoon committed
708
        trace!("---coloring done---");
709 710
        while !self.select_stack.is_empty() {
            let n = self.select_stack.pop().unwrap();
711
            trace!("Assigning color to {}", n);
qinsoon's avatar
qinsoon committed
712 713 714

            let mut ok_colors: LinkedHashSet<MuID> =
                self.colors.get(&self.ig.get_group_of(n)).unwrap().clone();
qinsoon's avatar
qinsoon committed
715 716 717

            trace!("all the colors for this temp: {:?}", ok_colors);

718 719
            for w in self.ig.get_adj_list(n).iter() {
                let w_alias = self.get_alias(*w);
qinsoon's avatar
qinsoon committed
720
                match self.ig.get_color_of(w_alias) {
qinsoon's avatar
qinsoon committed
721
                    None => {} // do nothing
qinsoon's avatar
qinsoon committed
722
                    Some(color) => {
qinsoon's avatar
qinsoon committed
723 724 725
                        trace!(
                            "color {} is used for its neighbor {:?} (aliasing to {:?})",
                            color,
726 727
                            w,
                            w_alias
qinsoon's avatar
qinsoon committed
728
                        );
qinsoon's avatar
qinsoon committed
729 730
                        ok_colors.remove(&color);
                    }
731 732
                }
            }
qinsoon's avatar
qinsoon committed
733
            trace!("available colors: {:?}", ok_colors);
qinsoon's avatar
qinsoon committed
734

735
            if ok_colors.is_empty() {
736
                trace!("{} is a spilled node", n);
737 738
                self.spilled_nodes.push(n);
            } else {
739
                let first_available_color = ok_colors.pop_front().unwrap();
740
                trace!("Color {} as {}", n, first_available_color);
qinsoon's avatar
qinsoon committed
741

742
                if !backend::is_callee_saved(first_available_color) {
743
                    trace!("Use caller saved register {}", first_available_color);
744
                }
qinsoon's avatar
qinsoon committed
745

746
                self.colored_nodes.push(n);
qinsoon's avatar
qinsoon committed
747
                self.ig.color_node(n, first_available_color);
748 749
            }
        }
qinsoon's avatar
qinsoon committed
750

qinsoon's avatar
qinsoon committed
751
        for n in self.coalesced_nodes.iter() {
752 753
            let n = *n;
            let alias = self.get_alias(n);
qinsoon's avatar
qinsoon committed
754
            if let Some(alias_color) = self.ig.get_color_of(alias) {
755 756
                trace!("Assign color to {} based on aliased {}", n, alias);
                trace!("Color {} as {}", n, alias_color);
qinsoon's avatar
qinsoon committed
757 758
                self.ig.color_node(n, alias_color);
            }
759 760
        }
    }
qinsoon's avatar
qinsoon committed
761

qinsoon's avatar
qinsoon committed
762
    fn rewrite_program(&mut self) {
qinsoon's avatar
qinsoon committed
763 764
        let spills = self.spills();

765
        let mut spilled_mem = LinkedHashMap::new();
qinsoon's avatar
qinsoon committed
766 767 768

        // allocating frame slots for every spilled temp
        for reg_id in spills.iter() {
qinsoon's avatar
qinsoon committed
769
            let ssa_entry = match self.func.context.get_value(*reg_id) {
qinsoon's avatar
qinsoon committed
770
                Some(entry) => entry,
771
                None => panic!("The spilled register {} is not in func context", reg_id)
qinsoon's avatar
qinsoon committed
772
            };
qinsoon's avatar
qinsoon committed
773 774 775
            let mem = self.cf
                .frame
                .alloc_slot_for_spilling(ssa_entry.value().clone(), self.vm);
qinsoon's avatar
qinsoon committed
776

qinsoon's avatar
qinsoon committed
777 778
            spilled_mem.insert(*reg_id, mem.clone());
            self.spill_history.insert(*reg_id, mem);
qinsoon's avatar
qinsoon committed
779 780
        }

qinsoon's avatar
qinsoon committed
781 782 783 784
        let scratch_temps = backend::spill_rewrite(&spilled_mem, self.func, self.cf, self.vm);
        for (k, v) in scratch_temps {
            self.spill_scratch_temps.insert(k, v);
        }
qinsoon's avatar
qinsoon committed
785
    }
qinsoon's avatar
qinsoon committed
786

787 788
    pub fn spills(&self) -> Vec<MuID> {
        let mut spills = vec![];
qinsoon's avatar
qinsoon committed
789

790 791 792 793 794 795
        let spill_count = self.spilled_nodes.len();
        if spill_count > 0 {
            for n in self.spilled_nodes.iter() {
                spills.push(self.ig.get_temp_of(*n));
            }
        }
qinsoon's avatar
qinsoon committed
796

797
        spills
798
    }
799 800 801 802 803 804 805 806 807 808 809 810 811

    pub fn get_assignments(&self) -> LinkedHashMap<MuID, MuID> {
        let mut ret = LinkedHashMap::new();

        for node in self.ig.nodes() {
            let temp = self.ig.get_temp_of(node);

            if temp < MACHINE_ID_END {
                continue;
            } else {
                let alias = self.get_alias(node);
                let machine_reg = match self.ig.get_color_of(alias) {
                    Some(reg) => reg,
qinsoon's avatar
qinsoon committed
812 813 814 815 816 817 818 819 820
                    None => {
                        panic!(
                            "Reg{}/{:?} (aliased as Reg{}/{:?}) is not assigned with a color",
                            self.ig.get_temp_of(node),
                            node,
                            self.ig.get_temp_of(alias),
                            alias
                        )
                    }
821 822 823 824 825
                };

                ret.insert(temp, machine_reg);
            }
        }
qinsoon's avatar
qinsoon committed
826

827 828 829
        ret
    }

qinsoon's avatar
qinsoon committed
830 831 832
    pub fn get_spill_scratch_temps(&self) -> LinkedHashMap<MuID, MuID> {
        self.spill_scratch_temps.clone()
    }
833
}