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 33.1 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
use utils::vec_utils;
25
use utils::LinkedHashSet;
26
use utils::LinkedHashMap;
27 28
use compiler::backend::reg_alloc::graph_coloring::liveness::Move;

29
/// allows coalescing
qinsoon's avatar
qinsoon committed
30
const COALESCING: bool = true;
31 32 33 34 35 36 37
/// abort after N rewrite iterations
/// (this is used to detect any possible infinite loop due to bugs)
const MAX_REWRITE_ITERATIONS_ALLOWED: usize = 50;
/// check invariants in every loop
/// (this will make register allocation run extremely slow - be careful
/// when using this with large workloads)
const CHECK_INVARIANTS: bool = false;
38

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

48 49 50 51
    /// 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
52
    /// machine registers, preassigned a color
53
    precolored: LinkedHashSet<MuID>,
qinsoon's avatar
qinsoon committed
54
    /// all colors available
55
    colors: LinkedHashMap<backend::RegGroup, LinkedHashSet<MuID>>,
qinsoon's avatar
qinsoon committed
56
    /// temporaries, not precolored and not yet processed
qinsoon's avatar
qinsoon committed
57
    initial: LinkedHashSet<MuID>,
qinsoon's avatar
qinsoon committed
58 59

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

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

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

qinsoon's avatar
qinsoon committed
93 94 95 96
    // 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
97
    spill_scratch_temps: LinkedHashMap<MuID, MuID>
qinsoon's avatar
qinsoon committed
98 99
}

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

qinsoon's avatar
qinsoon committed
117
    /// restarts coloring with spill history
qinsoon's avatar
qinsoon committed
118 119 120 121 122 123
    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,
124
        vm: &'a VM
qinsoon's avatar
qinsoon committed
125 126 127
    ) -> GraphColoring<'a> {
        assert!(
            iteration_count < MAX_REWRITE_ITERATIONS_ALLOWED,
128
            "reach graph coloring max rewrite iterations ({}), probably something is going wrong",
qinsoon's avatar
qinsoon committed
129 130
            MAX_REWRITE_ITERATIONS_ALLOWED
        );
131 132
        let iteration_count = iteration_count + 1;

133
        trace!("Initializing coloring allocator...");
qinsoon's avatar
qinsoon committed
134 135
        cf.mc().trace_mc();

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

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

qinsoon's avatar
qinsoon committed
171
        coloring.regalloc()
qinsoon's avatar
qinsoon committed
172
    }
qinsoon's avatar
qinsoon committed
173

qinsoon's avatar
qinsoon committed
174
    /// returns formatted string for an ID
qinsoon's avatar
qinsoon committed
175 176 177 178
    fn display_id(&self, id: MuID) -> String {
        self.func.context.get_temp_display(id)
    }

qinsoon's avatar
qinsoon committed
179
    /// returns formatted string for a move
qinsoon's avatar
qinsoon committed
180
    fn display_move(&self, m: Move) -> String {
qinsoon's avatar
qinsoon committed
181 182
        format!(
            "Move: {} -> {}",
183 184
            self.display_id(m.from),
            self.display_id(m.to)
qinsoon's avatar
qinsoon committed
185
        )
qinsoon's avatar
qinsoon committed
186
    }
qinsoon's avatar
qinsoon committed
187 188

    /// does coloring register allocation
189
    fn regalloc(mut self) -> GraphColoring<'a> {
qinsoon's avatar
qinsoon committed
190
        trace!("---InterenceGraph---");
qinsoon's avatar
qinsoon committed
191
        self.ig.print(&self.func.context);
qinsoon's avatar
qinsoon committed
192 193 194

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

qinsoon's avatar
qinsoon committed
196
        // precolor for all machine registers
197
        for reg in backend::all_regs().values() {
198
            let reg_id = reg.extract_ssa_id().unwrap();
199
            self.precolored.insert(reg_id);
200
        }
qinsoon's avatar
qinsoon committed
201

qinsoon's avatar
qinsoon committed
202
        // put usable registers as available colors
203
        for reg in backend::all_usable_regs().iter() {
qinsoon's avatar
qinsoon committed
204
            let reg_id = reg.extract_ssa_id().unwrap();
qinsoon's avatar
qinsoon committed
205 206
            let group = backend::pick_group_for_reg(reg_id);
            self.colors.get_mut(&group).unwrap().insert(reg_id);
qinsoon's avatar
qinsoon committed
207
        }
qinsoon's avatar
qinsoon committed
208 209

        // push uncolored nodes to initial work set
qinsoon's avatar
qinsoon committed
210 211
        for node in self.ig.nodes() {
            if !self.ig.is_colored(node) {
qinsoon's avatar
qinsoon committed
212
                self.initial.insert(node);
qinsoon's avatar
qinsoon committed
213 214
            }
        }
qinsoon's avatar
qinsoon committed
215 216

        // initialize work
qinsoon's avatar
qinsoon committed
217 218
        self.build();
        self.make_work_list();
qinsoon's avatar
qinsoon committed
219 220

        // main loop
qinsoon's avatar
qinsoon committed
221 222 223 224 225 226 227 228 229 230
        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
231

qinsoon's avatar
qinsoon committed
232 233 234 235
            if CHECK_INVARIANTS {
                self.check_invariants();
            }

qinsoon's avatar
qinsoon committed
236 237
            !(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
238
        } {}
qinsoon's avatar
qinsoon committed
239 240

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

qinsoon's avatar
qinsoon committed
243
        // finish
244 245
        drop(_p);

qinsoon's avatar
qinsoon committed
246
        // if we need to spill
qinsoon's avatar
qinsoon committed
247 248
        if !self.spilled_nodes.is_empty() {
            trace!("spill required");
qinsoon's avatar
qinsoon committed
249 250 251
            if cfg!(debug_assertions) {
                trace!("nodes to be spilled:");
                for node in self.spilled_nodes.iter() {
252
                    trace!("{}", *node);
qinsoon's avatar
qinsoon committed
253 254 255
                }
            }

qinsoon's avatar
qinsoon committed
256
            // rewrite program to insert spilling code
qinsoon's avatar
qinsoon committed
257
            self.rewrite_program();
qinsoon's avatar
qinsoon committed
258

qinsoon's avatar
qinsoon committed
259
            // recursively redo graph coloring
qinsoon's avatar
qinsoon committed
260 261 262 263 264 265
            return GraphColoring::start_with_spill_history(
                self.spill_history.clone(),
                self.spill_scratch_temps.clone(),
                self.iteration_count,
                self.func,
                self.cf,
266
                self.vm
qinsoon's avatar
qinsoon committed
267
            );
qinsoon's avatar
qinsoon committed
268 269
        }

270
        self
qinsoon's avatar
qinsoon committed
271
    }
qinsoon's avatar
qinsoon committed
272

qinsoon's avatar
qinsoon committed
273 274 275 276 277 278 279 280 281 282 283 284 285 286 287 288 289 290 291 292 293 294 295 296 297 298 299 300 301 302 303 304 305 306 307 308 309 310 311 312 313 314 315 316 317 318 319 320 321 322 323 324 325 326 327 328 329 330 331 332 333 334 335 336 337 338 339 340 341 342 343 344 345 346 347 348 349 350 351 352 353 354 355 356 357 358 359 360 361 362 363 364 365 366 367 368 369 370 371 372 373 374 375 376 377 378 379 380 381 382 383 384 385 386 387 388 389 390 391 392 393 394 395 396 397 398 399 400 401 402 403 404 405 406 407 408 409 410 411 412 413 414 415 416 417 418 419 420 421 422 423 424 425 426 427 428 429 430 431 432 433 434
    fn check_invariants(&self) {
        self.checkinv_degree();
        self.checkinv_simplify_worklist();
        self.checkinv_freeze_worklist();
        self.checkinv_spill_worklist();
    }

    fn checkinv_degree(&self) {
        trace!("Check invariant: degree...");
        // degree invariant
        // for u in simplifyWorklist \/ freezeWorklist \/ spillWorklist =>
        //   degree(u) = |adjList(u) /\
        //               (precolored \/ simplifyWorklist \/ freezeWorklist \/ spillWorklist)|

        let union = {
            let mut ret = LinkedHashSet::new();
            ret.add_all(self.worklist_simplify.clone());
            ret.add_all(self.worklist_freeze.clone());
            ret.add_all(self.worklist_spill.clone());
            ret
        };

        for u in union.iter() {
            let degree_u = self.ig.get_degree_of(*u);

            let set: LinkedHashSet<MuID> = {
                let adj: LinkedHashSet<MuID> = self.ig.get_adj_list(*u).clone();

                let mut union_: LinkedHashSet<MuID> = LinkedHashSet::new();
                union_.add_all(self.precolored.clone());
                union_.add_all(union.clone());

                {
                    let mut intersect: LinkedHashSet<MuID> = LinkedHashSet::new();
                    for item in adj.iter() {
                        if union_.contains(item) {
                            intersect.insert(*item);
                        }
                    }
                    intersect
                }
            };

            self.checkinv_assert(
                degree_u == set.len(),
                format!("degree({})={} != set(len)={}", u, degree_u, set.len())
            );
        }
    }

    fn checkinv_simplify_worklist(&self) {
        trace!("Check invariant: worklistSimplify...");
        // for u in simplifyWorkList
        //   either u is 'selected for spilling', otherwise
        //   degree(u) < K &&
        //   movelist(u) /\ (activeMoves \/ worklistMoves) = {} (emtpy)
        for u in self.worklist_simplify.iter() {
            if self.waiting_for_spill.contains(&u) {
                // no longer needs to check
                return;
            } else {
                // 1st cond: degree(u) < K
                let degree = self.ig.get_degree_of(*u);
                self.checkinv_assert(
                    degree < self.n_regs_for_node(*u),
                    format!("degree({})={} < K fails", u, degree)
                );

                // 2nd cond
                let movelist = self.get_movelist(*u);
                let union = {
                    let mut ret = self.active_moves.clone();
                    ret.add_all(self.worklist_moves.clone());
                    ret
                };
                let intersect = {
                    let mut ret = LinkedHashSet::new();
                    for m in movelist.iter() {
                        if union.contains(m) {
                            ret.insert(*m);
                        }
                    }
                    ret
                };

                self.checkinv_assert(
                    intersect.len() == 0,
                    format!("intersect({}) is not empty", u)
                );
            }
        }
    }

    fn checkinv_freeze_worklist(&self) {
        trace!("Check invariant: worklistFreeze...");
        // for u in freezeWorklist
        //   degree(u) < K &&
        //   moveList(u) /\ (activeMoves \/ worklistMoves) != {} (non empty)
        for u in self.worklist_freeze.iter() {
            // 1st cond: degree(u) < K
            let degree = self.ig.get_degree_of(*u);
            self.checkinv_assert(
                degree < self.n_regs_for_node(*u),
                format!("degree({})={} < K fails", u, degree)
            );

            // 2nd cond
            // 2nd cond
            let movelist = self.get_movelist(*u);
            let union = {
                let mut ret = self.active_moves.clone();
                ret.add_all(self.worklist_moves.clone());
                ret
            };
            let intersect = {
                let mut ret = LinkedHashSet::new();
                for m in movelist.iter() {
                    if union.contains(m) {
                        ret.insert(*m);
                    }
                }
                ret
            };
            self.checkinv_assert(intersect.len() != 0, format!("intersect({}) is empty", u));
        }
    }

    fn checkinv_spill_worklist(&self) {
        trace!("Check invariant: worklistSpill...");
        // for u in spillWorklist
        //    degree(u) >= K
        for u in self.worklist_spill.iter() {
            let degree = self.ig.get_degree_of(*u);
            self.checkinv_assert(
                degree >= self.n_regs_for_node(*u),
                format!("degree({})={} >= K fails", u, degree)
            );
        }
    }

    fn checkinv_assert(&self, cond: bool, msg: String) {
        if !cond {
            error!("{}", msg);

            // dump current state
            trace!("Current state:");

            trace!("simplifyWorklist: {:?}", self.worklist_simplify);
            trace!("freezeWorklist: {:?}", self.worklist_freeze);
            trace!("spillWorklist: {:?}", self.worklist_spill);
            trace!("worklistMoves: {:?}", self.worklist_moves);

            for node in self.ig.nodes() {
                trace!("Node {}: degree={}", node, self.ig.get_degree_of(node));
                trace!("         adjList={:?}", self.ig.get_adj_list(node));
                trace!("         moveList={:?}", self.get_movelist(node));
            }

            panic!()
        }
    }

435 436 437 438 439 440 441 442 443 444 445 446 447 448 449 450 451 452 453 454 455 456 457
    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
458
    fn build(&mut self) {
qinsoon's avatar
qinsoon committed
459
        if COALESCING {
qinsoon's avatar
qinsoon committed
460
            trace!("Coalescing enabled, build move list...");
qinsoon's avatar
qinsoon committed
461
            let ref ig = self.ig;
462
            for m in ig.moves().iter() {
qinsoon's avatar
qinsoon committed
463
                trace!("add to movelist: {:?}", m);
qinsoon's avatar
qinsoon committed
464
                self.worklist_moves.insert(*m);
465 466 467

                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
468
            }
qinsoon's avatar
qinsoon committed
469
        } else {
qinsoon's avatar
qinsoon committed
470
            trace!("Coalescing disabled...");
qinsoon's avatar
qinsoon committed
471 472
        }
    }
qinsoon's avatar
qinsoon committed
473

qinsoon's avatar
qinsoon committed
474
    fn make_work_list(&mut self) {
475
        trace!("Making work list from initials...");
qinsoon's avatar
qinsoon committed
476
        while !self.initial.is_empty() {
qinsoon's avatar
qinsoon committed
477
            let node = self.initial.pop_front().unwrap();
qinsoon's avatar
qinsoon committed
478

qinsoon's avatar
qinsoon committed
479 480
            // degree >= K
            if self.ig.get_degree_of(node) >= self.n_regs_for_node(node) {
qinsoon's avatar
qinsoon committed
481 482
                trace!(
                    "{} 's degree >= reg number limit (K), push to spill list",
483
                    node
qinsoon's avatar
qinsoon committed
484
                );
qinsoon's avatar
qinsoon committed
485
                self.worklist_spill.insert(node);
qinsoon's avatar
qinsoon committed
486
            } else if self.is_move_related(node) {
487
                trace!("{} is move related, push to freeze list", node);
qinsoon's avatar
qinsoon committed
488 489
                self.worklist_freeze.insert(node);
            } else {
qinsoon's avatar
qinsoon committed
490 491
                trace!(
                    "{} has small degree and not move related, push to simplify list",
492
                    node
qinsoon's avatar
qinsoon committed
493
                );
qinsoon's avatar
qinsoon committed
494 495 496 497
                self.worklist_simplify.insert(node);
            }
        }
    }
qinsoon's avatar
qinsoon committed
498

499
    fn n_regs_for_node(&self, node: MuID) -> usize {
500
        backend::number_of_usable_regs_in_group(self.ig.get_group_of(node))
501
    }
qinsoon's avatar
qinsoon committed
502

503
    fn is_move_related(&mut self, node: MuID) -> bool {
qinsoon's avatar
qinsoon committed
504 505
        !self.node_moves(node).is_empty()
    }
506 507 508 509 510 511 512 513 514 515

    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
516

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

qinsoon's avatar
qinsoon committed
520 521 522 523
        // addAll(active_moves)
        for m in self.active_moves.iter() {
            moves.insert(m.clone());
        }
qinsoon's avatar
qinsoon committed
524

qinsoon's avatar
qinsoon committed
525 526 527 528
        // addAll(worklist_moves)
        for m in self.worklist_moves.iter() {
            moves.insert(m.clone());
        }
qinsoon's avatar
qinsoon committed
529

530
        let mut retained = LinkedHashSet::new();
531
        let movelist = self.get_movelist(node);
qinsoon's avatar
qinsoon committed
532
        for m in moves.iter() {
533
            if movelist.contains(m) {
qinsoon's avatar
qinsoon committed
534 535 536
                retained.insert(*m);
            }
        }
qinsoon's avatar
qinsoon committed
537

qinsoon's avatar
qinsoon committed
538 539
        retained
    }
qinsoon's avatar
qinsoon committed
540

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

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

547
        self.select_stack.push(node);
548
        for m in self.adjacent(node).iter() {
549
            self.decrement_degree(*m);
550 551
        }
    }
qinsoon's avatar
qinsoon committed
552

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

556
        // add n's successors
557 558
        for s in self.ig.get_adj_list(n).iter() {
            adj.insert(*s);
559
        }
qinsoon's avatar
qinsoon committed
560

561 562 563 564
        // removeAll(select_stack)
        for s in self.select_stack.iter() {
            adj.remove(s);
        }
qinsoon's avatar
qinsoon committed
565

566 567 568 569
        // removeAll(coalesced_nodes)
        for s in self.coalesced_nodes.iter() {
            adj.remove(s);
        }
qinsoon's avatar
qinsoon committed
570

571 572
        adj
    }
qinsoon's avatar
qinsoon committed
573

574
    fn decrement_degree(&mut self, n: MuID) {
575 576 577
        if self.precolored.contains(&n) {
            return;
        }
qinsoon's avatar
qinsoon committed
578

579
        let d = self.ig.get_degree_of(n);
580
        debug_assert!(d != 0);
581
        self.ig.set_degree_of(n, d - 1);
qinsoon's avatar
qinsoon committed
582
        trace!("decrement degree of {} from {} to {}", n, d, d - 1);
qinsoon's avatar
qinsoon committed
583

584
        if d == self.n_regs_for_node(n) {
585
            trace!("{}'s degree is K, no longer need to spill it", n);
586 587 588
            let mut nodes = self.adjacent(n);
            nodes.insert(n);
            self.enable_moves(nodes);
qinsoon's avatar
qinsoon committed
589

qinsoon's avatar
qinsoon committed
590
            self.worklist_spill.remove(&n);
qinsoon's avatar
qinsoon committed
591

592
            if self.is_move_related(n) {
593
                trace!("{} is move related, push to freeze list", n);
594 595
                self.worklist_freeze.insert(n);
            } else {
596
                trace!("{} is not move related, push to simplify list", n);
597 598 599 600
                self.worklist_simplify.insert(n);
            }
        }
    }
qinsoon's avatar
qinsoon committed
601

602
    fn enable_moves(&mut self, nodes: LinkedHashSet<MuID>) {
qinsoon's avatar
qinsoon committed
603
        trace!("enable moves of: {:?}", nodes);
604 605 606 607
        for n in nodes.iter() {
            let n = *n;
            for mov in self.node_moves(n).iter() {
                let mov = *mov;
608
                if self.active_moves.contains(&mov) {
qinsoon's avatar
qinsoon committed
609
                    trace!("move {:?} from activeMoves to worklistMoves", mov);
qinsoon's avatar
qinsoon committed
610
                    self.active_moves.remove(&mov);
qinsoon's avatar
qinsoon committed
611
                    self.worklist_moves.insert(mov);
612 613 614 615
                }
            }
        }
    }
qinsoon's avatar
qinsoon committed
616

617
    fn coalesce(&mut self) {
qinsoon's avatar
qinsoon committed
618
        let m = self.worklist_moves.pop_front().unwrap();
qinsoon's avatar
qinsoon committed
619

qinsoon's avatar
qinsoon committed
620
        trace!("Coalescing on {:?}...", m);
qinsoon's avatar
qinsoon committed
621

qinsoon's avatar
qinsoon committed
622 623
        let x = self.get_alias(m.from);
        let y = self.get_alias(m.to);
qinsoon's avatar
qinsoon committed
624 625
        trace!("resolve alias: {} -> {}", m.from, x);
        trace!("resolve alias: {} -> {}", m.to, y);
qinsoon's avatar
qinsoon committed
626

qinsoon's avatar
qinsoon committed
627 628 629 630 631 632
        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
633

qinsoon's avatar
qinsoon committed
634 635 636 637 638 639
                (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
640

qinsoon's avatar
qinsoon committed
641 642 643
                (u, v, precolored_u, precolored_v)
            }
        };
qinsoon's avatar
qinsoon committed
644 645
        trace!(
            "u={}, v={}, precolored_u={}, precolroed_v={}",
646 647
            u,
            v,
qinsoon's avatar
qinsoon committed
648 649 650 651
            precolored_u,
            precolored_v
        );

qinsoon's avatar
qinsoon committed
652 653 654 655 656 657 658 659 660 661 662 663 664 665 666 667
        // if they are not from the same register group, we cannot coalesce them
        if self.ig.get_group_of(u) != self.ig.get_group_of(v) {
            if !precolored_v {
                self.add_worklist(v);
            }
            if !precolored_u {
                self.add_worklist(u);
            }
            self.constrained_moves.insert(m);
            info!(
                "u and v are temporaries of different register groups, cannot coalesce: {:?}",
                m
            );
            return;
        }

668
        // if u or v is a machine register that is not usable/not a color, we cannot coalesce
qinsoon's avatar
qinsoon committed
669 670 671 672 673 674 675 676
        if precolored_u {
            let group = backend::pick_group_for_reg(u);
            if !self.colors.get(&group).unwrap().contains(&u) {
                if !precolored_v {
                    self.add_worklist(v);
                }
                self.constrained_moves.insert(m);
                trace!("u is precolored but not a usable color, cannot coalesce");
677 678 679
                return;
            }
        }
qinsoon's avatar
qinsoon committed
680 681 682 683 684 685 686 687
        if precolored_v {
            let group = backend::pick_group_for_reg(v);
            if !self.colors.get(&group).unwrap().contains(&v) {
                if !precolored_u {
                    self.add_worklist(u);
                }
                self.constrained_moves.insert(m);
                trace!("v is precolored but not a usable color, cannot coalesce");
688 689 690 691
                return;
            }
        }

qinsoon's avatar
qinsoon committed
692
        if u == v {
qinsoon's avatar
qinsoon committed
693
            trace!("u == v, coalesce the move");
qinsoon's avatar
qinsoon committed
694 695 696 697
            self.coalesced_moves.insert(m);
            if !precolored_u {
                self.add_worklist(u);
            }
698
        } else if precolored_v || self.ig.is_in_adj_set(u, v) {
699
            trace!("precolored_v: {}", precolored_v);
700
            trace!("is_adj(u, v): {}", self.ig.is_in_adj_set(u, v));
qinsoon's avatar
qinsoon committed
701 702
            trace!("v is precolored or u,v is adjacent, the move is constrained");
            self.constrained_moves.insert(m);
qinsoon's avatar
qinsoon committed
703 704 705 706 707 708
            if !precolored_u {
                self.add_worklist(u);
            }
            if !precolored_v {
                self.add_worklist(v);
            }
709 710 711 712 713
        } else if (precolored_u && self.check_ok(u, v)) ||
                   (!precolored_u && self.check_conservative(u, v))
        {
            trace!("ok(u, v) = {}", self.check_ok(u, v));
            trace!("conservative(u, v) = {}", self.check_conservative(u, v));
qinsoon's avatar
qinsoon committed
714

qinsoon's avatar
qinsoon committed
715 716 717 718
            trace!(
                "precolored_u&&ok(u,v) || !precolored_u&&conserv(u,v), \
                 coalesce and combine the move"
            );
qinsoon's avatar
qinsoon committed
719 720 721 722 723 724
            self.coalesced_moves.insert(m);
            self.combine(u, v);
            if !precolored_u {
                self.add_worklist(u);
            }
        } else {
qinsoon's avatar
qinsoon committed
725
            trace!("cannot coalesce the move");
qinsoon's avatar
qinsoon committed
726 727 728
            self.active_moves.insert(m);
        }
    }
qinsoon's avatar
qinsoon committed
729

730
    pub fn get_alias(&self, node: MuID) -> MuID {
qinsoon's avatar
qinsoon committed
731 732 733 734 735 736
        if self.coalesced_nodes.contains(&node) {
            self.get_alias(*self.alias.get(&node).unwrap())
        } else {
            node
        }
    }
qinsoon's avatar
qinsoon committed
737

738 739
    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
740 741 742 743
            self.worklist_freeze.remove(&node);
            self.worklist_simplify.insert(node);
        }
    }
qinsoon's avatar
qinsoon committed
744

745
    fn check_ok(&self, u: MuID, v: MuID) -> bool {
746 747
        for t in self.adjacent(v).iter() {
            let t = *t;
748
            if !self.ok(t, u) {
qinsoon's avatar
qinsoon committed
749
                return false;
750
            }
qinsoon's avatar
qinsoon committed
751
        }
qinsoon's avatar
qinsoon committed
752

qinsoon's avatar
qinsoon committed
753 754
        true
    }
qinsoon's avatar
qinsoon committed
755

756 757 758 759 760 761 762 763
    fn ok(&self, t: MuID, r: MuID) -> bool {
        let degree_t = self.ig.get_degree_of(t);
        let k = self.n_regs_for_node(t);

        degree_t < k || self.precolored.contains(&t) || self.ig.is_in_adj_set(t, r)
    }

    fn check_conservative(&self, u: MuID, v: MuID) -> bool {
qinsoon's avatar
qinsoon committed
764 765
        let adj_u = self.adjacent(u);
        let adj_v = self.adjacent(v);
766 767 768 769 770
        let nodes = {
            let mut ret = adj_u;
            ret.add_all(adj_v);
            ret
        };
qinsoon's avatar
qinsoon committed
771

772 773 774
        let n_regs_for_group = self.n_regs_for_node(u);
        self.conservative(nodes, n_regs_for_group)
    }
qinsoon's avatar
qinsoon committed
775

776
    fn conservative(&self, nodes: LinkedHashSet<MuID>, n_regs_for_group: usize) -> bool {
qinsoon's avatar
qinsoon committed
777
        let mut k = 0;
778
        for n in nodes.iter() {
779 780
            // TODO: do we check if n is precolored?
            if self.precolored.contains(n) || self.ig.get_degree_of(*n) >= n_regs_for_group {
qinsoon's avatar
qinsoon committed
781 782 783
                k += 1;
            }
        }
784
        k < n_regs_for_group
qinsoon's avatar
qinsoon committed
785
    }
qinsoon's avatar
qinsoon committed
786

787
    fn combine(&mut self, u: MuID, v: MuID) {
qinsoon's avatar
qinsoon committed
788
        trace!("Combine temps {} and {}...", u, v);
qinsoon's avatar
qinsoon committed
789
        if self.worklist_freeze.contains(&v) {
qinsoon's avatar
qinsoon committed
790
            trace!("remove {} from freeze list", v);
qinsoon's avatar
qinsoon committed
791 792
            self.worklist_freeze.remove(&v);
        } else {
qinsoon's avatar
qinsoon committed
793 794
            trace!("remove {} from spill list", v);
            self.worklist_spill.remove(&v);
qinsoon's avatar
qinsoon committed
795
        }
qinsoon's avatar
qinsoon committed
796
        self.coalesced_nodes.insert(v);
qinsoon's avatar
qinsoon committed
797 798 799

        self.alias.insert(v, u);

qinsoon's avatar
qinsoon committed
800
        {
801 802
            // movelist[u] <- movelist[u] + movelist[v]
            let movelist_v = self.get_movelist(v);
qinsoon's avatar
qinsoon committed
803

804 805 806
            for m in movelist_v.iter() {
                GraphColoring::add_to_movelist(&mut self.movelist, u, *m)
            }
qinsoon's avatar
qinsoon committed
807
        }
qinsoon's avatar
qinsoon committed
808

809
        let mut nodes = LinkedHashSet::new();
qinsoon's avatar
qinsoon committed
810 811
        nodes.insert(v);
        self.enable_moves(nodes);
qinsoon's avatar
qinsoon committed
812

813 814
        for t in self.adjacent(v).iter() {
            let t = *t;
qinsoon's avatar
qinsoon committed
815 816 817
            self.add_edge(t, u);
            self.decrement_degree(t);
        }
qinsoon's avatar
qinsoon committed
818

819 820 821
        if self.worklist_freeze.contains(&u) &&
            self.ig.get_degree_of(u) >= self.n_regs_for_node(u)
        {
qinsoon's avatar
qinsoon committed
822
            self.worklist_freeze.remove(&u);
qinsoon's avatar
qinsoon committed
823
            self.worklist_spill.insert(u);
qinsoon's avatar
qinsoon committed
824 825
        }
    }
qinsoon's avatar
qinsoon committed
826

827
    fn add_edge(&mut self, u: MuID, v: MuID) {
qinsoon's avatar
qinsoon committed
828
        self.ig.add_edge(u, v);
829
    }
qinsoon's avatar
qinsoon committed
830

831
    fn freeze(&mut self) {
qinsoon's avatar
qinsoon committed
832
        // it is not empty (checked before)
833
        let node = self.worklist_freeze.pop_front().unwrap();
834
        trace!("Freezing {}...", node);
qinsoon's avatar
qinsoon committed
835

qinsoon's avatar
qinsoon committed
836 837 838
        self.worklist_simplify.insert(node);
        self.freeze_moves(node);
    }
qinsoon's avatar
qinsoon committed
839

840
    fn freeze_moves(&mut self, u: MuID) {
841 842
        for m in self.node_moves(u).iter() {
            let m = *m;
qinsoon's avatar
qinsoon committed
843 844 845 846
            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
847

qinsoon's avatar
qinsoon committed
848 849
            self.active_moves.remove(&m);
            self.frozen_moves.insert(m);
qinsoon's avatar
qinsoon committed
850 851

            if !self.precolored.contains(&v) && self.node_moves(v).is_empty() &&
852
                self.ig.get_degree_of(v) < self.n_regs_for_node(v)
qinsoon's avatar
qinsoon committed
853
            {
qinsoon's avatar
qinsoon committed
854 855 856 857
                self.worklist_freeze.remove(&v);
                self.worklist_simplify.insert(v);
            }
        }
858
    }
qinsoon's avatar
qinsoon committed
859

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

864 865 866 867 868
        for n in self.worklist_spill.iter() {
            let n = *n;
            if m.is_none() {
                m = Some(n);
            } else if {
qinsoon's avatar
qinsoon committed
869 870 871 872
                       // m is not none
                       let temp = self.ig.get_temp_of(m.unwrap());
                       !self.is_spillable(temp)
                   } {
873
                m = Some(n);
874 875 876
            } 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
877
            {
878 879 880
                m = Some(n);
            }
        }
qinsoon's avatar
qinsoon committed
881

882 883
        // m is not none
        let m = m.unwrap();
884
        trace!("Spilling {}...", m);
qinsoon's avatar
qinsoon committed
885

qinsoon's avatar
qinsoon committed
886 887
        self.waiting_for_spill.insert(m);
        self.worklist_spill.remove(&m);
888 889
        self.worklist_simplify.insert(m);
        self.freeze_moves(m);
890
    }
qinsoon's avatar
qinsoon committed
891

892
    fn assign_colors(&mut self) {
qinsoon's avatar
qinsoon committed
893
        trace!("---coloring done---");
894 895
        while !self.select_stack.is_empty() {
            let n = self.select_stack.pop().unwrap();
896
            trace!("Assigning color to {}", n);
qinsoon's avatar
qinsoon committed
897 898 899

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

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

903 904
            for w in self.ig.get_adj_list(n).iter() {
                let w_alias = self.get_alias(*w);
qinsoon's avatar
qinsoon committed
905
                match self.ig.get_color_of(w_alias) {
qinsoon's avatar
qinsoon committed
906
                    None => {} // do nothing
qinsoon's avatar
qinsoon committed
907
                    Some(color) => {
qinsoon's avatar
qinsoon committed
908 909 910
                        trace!(
                            "color {} is used for its neighbor {:?} (aliasing to {:?})",
                            color,
911 912
                            w,
                            w_alias
qinsoon's avatar
qinsoon committed
913
                        );
qinsoon's avatar
qinsoon committed
914 915
                        ok_colors.remove(&color);
                    }
916 917
                }
            }
qinsoon's avatar
qinsoon committed
918
            trace!("available colors: {:?}", ok_colors);
qinsoon's avatar
qinsoon committed
919

920
            if ok_colors.is_empty() {
921
                trace!("{} is a spilled node", n);
qinsoon's avatar
qinsoon committed
922
                self.spilled_nodes.insert(n);
923
            } else {
924
                let first_available_color = ok_colors.pop_front().unwrap();
925
                trace!("Color {} as {}", n, first_available_color);
qinsoon's avatar
qinsoon committed
926

927
                if !backend::is_callee_saved(first_available_color) {
928
                    trace!("Use caller saved register {}", first_available_color);
929
                }
qinsoon's avatar
qinsoon committed
930

qinsoon's avatar
qinsoon committed
931
                self.colored_nodes.insert(n);
qinsoon's avatar
qinsoon committed
932
                self.ig.color_node(n, first_available_color);
933 934
            }
        }
qinsoon's avatar
qinsoon committed
935

qinsoon's avatar
qinsoon committed
936
        for n in self.coalesced_nodes.iter() {
937 938
            let n = *n;
            let alias = self.get_alias(n);
qinsoon's avatar
qinsoon committed
939
            if let Some(alias_color) = self.ig.get_color_of(alias) {
940 941
                trace!("Assign color to {} based on aliased {}", n, alias);
                trace!("Color {} as {}", n, alias_color);
qinsoon's avatar
qinsoon committed
942 943
                self.ig.color_node(n, alias_color);
            }
944 945
        }
    }
qinsoon's avatar
qinsoon committed
946

qinsoon's avatar
qinsoon committed
947
    fn rewrite_program(&mut self) {
qinsoon's avatar
qinsoon committed
948 949
        let spills = self.spills();

950
        let mut spilled_mem = LinkedHashMap::new();
qinsoon's avatar
qinsoon committed
951 952 953

        // allocating frame slots for every spilled temp
        for reg_id in spills.iter() {
qinsoon's avatar
qinsoon committed
954
            let ssa_entry = match self.func.context.get_value(*reg_id) {
qinsoon's avatar
qinsoon committed
955
                Some(entry) => entry,
956
                None => panic!("The spilled register {} is not in func context", reg_id)
qinsoon's avatar
qinsoon committed
957
            };
qinsoon's avatar
qinsoon committed
958 959 960
            let mem = self.cf
                .frame
                .alloc_slot_for_spilling(ssa_entry.value().clone(), self.vm);
qinsoon's avatar
qinsoon committed
961

qinsoon's avatar
qinsoon committed
962 963
            spilled_mem.insert(*reg_id, mem.clone());
            self.spill_history.insert(*reg_id, mem);
qinsoon's avatar
qinsoon committed
964 965
        }

qinsoon's avatar
qinsoon committed
966 967 968 969
        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
970
    }
qinsoon's avatar
qinsoon committed
971

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

975 976 977 978 979 980
        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
981

982
        spills
983
    }
984 985 986 987 988 989 990 991 992 993 994 995 996

    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
997 998 999 1000 1001 1002 1003 1004 1005
                    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
                        )
                    }
1006 1007 1008 1009 1010
                };

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

1012 1013 1014
        ret
    }

qinsoon's avatar
qinsoon committed
1015 1016 1017
    pub fn get_spill_scratch_temps(&self) -> LinkedHashMap<MuID, MuID> {
        self.spill_scratch_temps.clone()
    }
1018
}