<commit_before>use std::cmp::{max, Reverse};
use std::fmt;
use std::collections::{BTreeMap, BTreeSet, BinaryHeap};
use std::convert::TryFrom;

use derive_more::{Add, From, LowerHex, Sub};
use num_enum::TryFromPrimitive;

use crate::serialize::Record;

const TASK_GRANULARITY_THRESHOLD: u64 = 10 * 1000;


// Make sure this is up to date with lowlevel.h
#[derive(Debug, Copy, Clone, Eq, PartialEq, Ord, PartialOrd, TryFromPrimitive)]
#[repr(i32)]
pub enum ProcKind {
    GPU = 1,
    CPU = 2,
    Utility = 3,
    IO = 4,
    ProcGroup = 5,
    OpenMP = 6,
    Python = 7,
}

// Make sure this is up to date with lowlevel.h
#[derive(Debug, Copy, Clone, Eq, PartialEq, TryFromPrimitive)]
#[repr(i32)]
pub enum MemKind {
    NoMemKind = 0,
    Global = 1,
    System = 2,
    Registered = 3,
    Socket = 4,
    ZeroCopy = 5,
    Framebuffer = 6,
    Disk = 7,
    HDF5 = 8,
    File = 9,
    L3Cache = 10,
    L2Cache = 11,
    L1Cache = 12,
}

// Make sure this is up to date with lowlevel.h
#[derive(Debug, Copy, Clone, Eq, PartialEq, TryFromPrimitive)]
#[repr(i32)]
pub enum DepPartKind {
    Union = 0,
    Unions = 1,
    UnionReduction = 2,
    Intersection = 3,
    Intersections = 4,
    IntersectionReduction = 5,
    Difference = 6,
    Differences = 7,
    EqualPartition = 8,
    PartitionByField = 9,
    PartitionByImage = 10,
    PartitionByImageRange = 11,
    PartitionByPreimage = 12,
    PartitionByPreimageRange = 13,
    CreateAssociation = 14,
    PartitionByWeights = 15,
}

// Make sure this is up to date with lowlevel.h
#[derive(Debug, Copy, Clone, Eq, PartialEq, TryFromPrimitive)]
#[repr(u32)]
pub enum DimKind {
    DimX = 0,
    DimY = 1,
    DimZ = 2,
    DimW = 3,
    DimV = 4,
    DimU = 5,
    DimT = 6,
    DimS = 7,
    DimR = 8,
    DimF = 9,
    InnerDimX = 10,
    OuterDimX = 11,
    InnerDimY = 12,
    OuterDimY = 13,
    InnerDimZ = 14,
    OuterDimZ = 15,
    InnerDimW = 16,
    OuterDimW = 17,
    InnerDimV = 18,
    OuterDimV = 19,
    InnerDimU = 20,
    OuterDimU = 21,
    InnerDimT = 22,
    OuterDimT = 23,
    InnerDimS = 24,
    OuterDimS = 25,
    InnerDimR = 26,
    OuterDimR = 27,
}

#[derive(Debug, Copy, Clone, PartialEq, Eq, PartialOrd, Ord, Default, Add, Sub, From)]
pub struct Timestamp(pub u64 /* ns */);

impl fmt::Display for Timestamp {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        // Time is stored in nanoseconds. But it is displayed in microseconds.
        let nanoseconds = self.0;
        let divisor = 1000;
        let microseconds = nanoseconds / divisor;
        let remainder = nanoseconds % divisor;
        write!(f, "{}.{:0>3}", microseconds, remainder)
    }
}

#[derive(Debug, Copy, Clone)]
pub enum ProcEntry {
    Task(OpID),
    MetaTask(OpID, VariantID, usize),
    MapperCall(usize),
    RuntimeCall(usize),
    ProfTask(usize)
}

#[derive(Debug, Copy, Clone)]
pub struct ProcPoint {
    pub time: Timestamp,
    pub entry: ProcEntry,
    pub first: bool,
}

impl ProcPoint {
    fn new(time: Timestamp, entry: ProcEntry, first: bool) -> Self {
        ProcPoint { time, entry, first }
    }
    pub fn time_key(&self) -> Timestamp {
        Timestamp(2 * self.time.0) + (1000 * (!self.first) as u64).into()
    }
}

#[derive(Debug, Copy, Clone, PartialEq, Eq, PartialOrd, Ord, LowerHex)]
pub struct ProcID(pub u64);

#[derive(Debug, Copy, Clone, PartialEq, Eq, PartialOrd, Ord)]
pub struct NodeID(pub u64);

#[derive(Debug, Copy, Clone, PartialEq, Eq, PartialOrd, Ord)]
pub struct EventID(pub u64);

#[derive(Debug)]
pub struct Proc {
    pub proc_id: ProcID,
    pub kind: ProcKind,
    pub node_id: NodeID,
    pub proc_in_node: u64,
    app_ranges: Vec<()>,
    last_time: Option<Timestamp>,
    pub tasks: BTreeMap<OpID, Task>,
    pub meta_tasks: BTreeMap<(OpID, VariantID), Vec<MetaTask>>,
    pub mapper_calls: Vec<MapperCall>,
    pub runtime_calls: Vec<RuntimeCall>,
    pub prof_tasks: Vec<ProfTask>,
    pub max_levels: u32,
    pub max_levels_ready: u32,
    pub time_points: Vec<ProcPoint>,
    pub util_time_points: Vec<ProcPoint>,
}

impl Proc {
    fn new(proc_id: ProcID, kind: ProcKind) -> Self {
        Proc {
            proc_id,
            kind,
            // PROCESSOR:   tag:8 = 0x1d, owner_node:16,   (unused):28, proc_idx: 12
            // owner_node = proc_id[55:40]
            // proc_idx = proc_id[11:0]
            node_id: NodeID((proc_id.0 >> 40) & ((1 << 16) - 1)),
            proc_in_node: (proc_id.0) & ((1 << 12) - 1),
            app_ranges: Vec::new(),
            last_time: None,
            tasks: BTreeMap::new(),
            meta_tasks: BTreeMap::new(),
            mapper_calls: Vec::new(),
            runtime_calls: Vec::new(),
            prof_tasks: Vec::new(),
            max_levels: 0,
            max_levels_ready: 0,
            time_points: Vec::new(),
            util_time_points: Vec::new(),
        }
    }

    pub fn is_empty(&self) -> bool {
        self.tasks.is_empty() && self.mapper_calls.is_empty() && self.runtime_calls.is_empty() && self.prof_tasks.is_empty()
    }

    pub fn entry(&self, entry: ProcEntry) -> (&Base, &TimeRange, &Waiters) {
        match entry {
            ProcEntry::Task(op_id) => {
                let task = self.tasks.get(&op_id).unwrap();
                (&task.base, &task.time_range, &task.waiters)
            }
            ProcEntry::MetaTask(op_id, variant_id, idx) => {
                let task = &self.meta_tasks.get(&(op_id, variant_id)).unwrap()[idx];
                (&task.base, &task.time_range, &task.waiters)
            }
            ProcEntry::MapperCall(idx) => {
                let call = &self.mapper_calls[idx];
                (&call.base, &call.time_range, &call.waiters)
            }
            ProcEntry::RuntimeCall(idx) => {
                let call = &self.runtime_calls[idx];
                (&call.base, &call.time_range, &call.waiters)
            }
            ProcEntry::ProfTask(idx) => {
                let task = &self.prof_tasks.get(idx).unwrap();
                (&task.base, &task.time_range, &task.waiters)
            }
        }
    }

    pub fn entry_mut(&mut self, entry: ProcEntry) -> (&mut Base, &mut TimeRange, &mut Waiters) {
        match entry {
            ProcEntry::Task(op_id) => {
                let task = self.tasks.get_mut(&op_id).unwrap();
                (&mut task.base, &mut task.time_range, &mut task.waiters)
            }
            ProcEntry::MetaTask(op_id, variant_id, idx) => {
                let task = &mut self.meta_tasks.get_mut(&(op_id, variant_id)).unwrap()[idx];
                (&mut task.base, &mut task.time_range, &mut task.waiters)
            }
            ProcEntry::MapperCall(idx) => {
                let call = &mut self.mapper_calls[idx];
                (&mut call.base, &mut call.time_range, &mut call.waiters)
            }
            ProcEntry::RuntimeCall(idx) => {
                let call = &mut self.runtime_calls[idx];
                (&mut call.base, &mut call.time_range, &mut call.waiters)
            }
            ProcEntry::ProfTask(idx) => {
                let task = self.prof_tasks.get_mut(idx).unwrap();
                (&mut task.base, &mut task.time_range, &mut task.waiters)
            }
        }
    }

    pub fn entry_base(&self, entry: ProcEntry) -> &Base {
        self.entry(entry).0
    }

    pub fn entry_base_mut(&mut self, entry: ProcEntry) -> &mut Base {
        self.entry_mut(entry).0
    }

    fn trim_time_range(&mut self, start: Timestamp, stop: Timestamp) {
        let mut removed_tasks = Vec::new();
        for (op_id, task) in self.tasks.iter_mut() {
            task.trim_time_range(start, stop);
            if task.time_range.was_removed {
                removed_tasks.push(*op_id);
            }
        }
        for op_id in removed_tasks {
            self.tasks.remove(&op_id);
        }
        for mapper_call in &mut self.mapper_calls {
            mapper_call.trim_time_range(start, stop);
        }
        self.mapper_calls.retain(|t| !t.time_range.was_removed);
        for runtime_call in &mut self.runtime_calls {
            runtime_call.trim_time_range(start, stop);
        }
        self.runtime_calls.retain(|t| !t.time_range.was_removed);
    }

    fn sort_time_range(&mut self, last_time: Timestamp) {
        fn add(
            time: &TimeRange,
            entry: ProcEntry,
            all_points: &mut Vec<ProcPoint>,
            points: &mut Vec<ProcPoint>,
            util_points: &mut Vec<ProcPoint>,
        ) {
            let start = time.start.unwrap();
            let stop = time.stop.unwrap();
            let ready = time.ready;
            if stop - start > TASK_GRANULARITY_THRESHOLD.into() && !ready.is_none() {
                all_points.push(ProcPoint::new(ready.unwrap(), entry, true));
                all_points.push(ProcPoint::new(stop, entry, false));
            } else {
                all_points.push(ProcPoint::new(start, entry, true));
                all_points.push(ProcPoint::new(stop, entry, false));
            }

            points.push(ProcPoint::new(start, entry, true));
            points.push(ProcPoint::new(stop, entry, false));

            util_points.push(ProcPoint::new(start, entry, true));
            util_points.push(ProcPoint::new(stop, entry, false));
        }
        fn add_waiters(waiters: &Waiters, entry: ProcEntry, util_points: &mut Vec<ProcPoint>) {
            for wait in &waiters.wait_intervals {
                util_points.push(ProcPoint::new(wait.start, entry, false));
                util_points.push(ProcPoint::new(wait.end, entry, true));
            }
        }

        let mut all_points = Vec::new();
        let mut points = Vec::new();
        let mut util_points = Vec::new();

        for task in self.tasks.values() {
            let time = &task.time_range;
            let entry = ProcEntry::Task(task.op_id);
            add(&time, entry, &mut all_points, &mut points, &mut util_points);
            add_waiters(&task.waiters, entry, &mut util_points);
        }
        for tasks in self.meta_tasks.values() {
            for (idx, task) in tasks.iter().enumerate() {
                let time = &task.time_range;
                let entry = ProcEntry::MetaTask(task.op_id, task.variant_id, idx);
                add(&time, entry, &mut all_points, &mut points, &mut util_points);
                add_waiters(&task.waiters, entry, &mut util_points);
            }
        }
        for (idx, mapper_call) in self.mapper_calls.iter().enumerate() {
            let time = &mapper_call.time_range;
            let entry = ProcEntry::MapperCall(idx);
            add(&time, entry, &mut all_points, &mut points, &mut util_points);
            add_waiters(&mapper_call.waiters, entry, &mut util_points);
        }
        for (idx, runtime_call) in self.runtime_calls.iter().enumerate() {
            let time = &runtime_call.time_range;
            let entry = ProcEntry::RuntimeCall(idx);
            add(&time, entry, &mut all_points, &mut points, &mut util_points);
            add_waiters(&runtime_call.waiters, entry, &mut util_points);
        }

        for (idx, prof_task) in self.prof_tasks.iter().enumerate() {
            let time = &prof_task.time_range;
            let entry = ProcEntry::ProfTask(idx);
            add(&time, entry, &mut all_points, &mut points, &mut util_points);
            add_waiters(&prof_task.waiters, entry, &mut util_points);
        }

        points.sort_by(|a, b| a.time_key().cmp(&b.time_key()));
        util_points.sort_by(|a, b| a.time_key().cmp(&b.time_key()));

        // Hack: This is a max heap so reverse the values as they go in.
        let mut free_levels = BinaryHeap::<Reverse<u32>>::new();
        for point in &points {
            if point.first {
                let level = if let Some(level) = free_levels.pop() {
                    level.0
                } else {
                    self.max_levels += 1;
                    self.max_levels
                };
                self.entry_base_mut(point.entry).set_level(level);
            } else {
                let level = self.entry_base(point.entry).level.unwrap();
                free_levels.push(Reverse(level));
            }
        }

        all_points.sort_by(|a, b| a.time_key().cmp(&b.time_key()));

        // Hack: This is a max heap so reverse the values as they go in.
        let mut free_levels_ready = BinaryHeap::<Reverse<u32>>::new();
        for point in &all_points {
            if point.first {
                let level = if let Some(level) = free_levels_ready.pop() {
                    level.0
                } else {
                    self.max_levels_ready += 1;
                    self.max_levels_ready
                };
                self.entry_base_mut(point.entry).set_level_ready(level);
            } else {
                let level = self.entry_base(point.entry).level_ready.unwrap();
                free_levels_ready.push(Reverse(level));
            }
        }

        self.time_points = points;
        self.util_time_points = util_points;
    }
}

#[derive(Debug, Copy, Clone, PartialEq, Eq, PartialOrd, Ord)]
pub struct MemID(pub u64);

impl MemID {
    pub fn node_id(&self) -> NodeID {
        NodeID((self.0 >> 40) & ((1 << 16) - 1))
    }
}

#[derive(Debug)]
pub struct Mem {
    mem_id: MemID,
    kind: MemKind,
    capacity: u64,
    node_id: NodeID,
    mem_in_node: u64,
    instances: BTreeSet<()>,
    time_points: Vec<()>,
    max_live_instances: Option<u64>,
    last_time: Option<Timestamp>,
    affinity: Option<u64>,
}

impl Mem {
    fn new(mem_id: MemID, kind: MemKind, capacity: u64) -> Self {
        Mem {
            mem_id,
            kind,
            capacity,
            node_id: NodeID((mem_id.0 >> 40) & ((1 << 16) - 1)),
            mem_in_node: (mem_id.0) & ((1 << 12) - 1),
            instances: BTreeSet::new(),
            time_points: Vec::new(),
            max_live_instances: None,
            last_time: None,
            affinity: None,
        }
    }
    fn trim_time_range(&mut self, start: Timestamp, stop: Timestamp) {}
    fn sort_time_range(&mut self, last_time: Timestamp) {}
}

#[derive(Debug)]
pub struct MemProcAffinity {
    mem_id: MemID,
    proc_ids: Vec<ProcID>,
}

impl MemProcAffinity {
    fn new(mem_id: MemID) -> Self {
        MemProcAffinity {
            mem_id,
            proc_ids: Vec::new(),
        }
    }
    fn add_proc_id(&mut self, proc_id: ProcID) {
        self.proc_ids.push(proc_id);
    }
}

#[derive(Debug, Copy, Clone)]
pub enum ChanEntry {
    Copy(usize),
    Fill(usize),
    DepPart(usize),
}

#[derive(Debug, Copy, Clone)]
pub struct ChanPoint {
    pub time: Timestamp,
    pub entry: ChanEntry,
    pub first: bool,
}

impl ChanPoint {
    fn new(time: Timestamp, entry: ChanEntry, first: bool) -> Self {
        ChanPoint { time, entry, first }
    }
    pub fn time_key(&self) -> Timestamp {
        Timestamp(2 * self.time.0) + (1000 * (!self.first) as u64).into()
    }
}

#[derive(Debug, Copy, Clone, PartialEq, Eq, PartialOrd, Ord)]
pub struct ChanID {
    pub src: Option<MemID>,
    pub dst: Option<MemID>,
}

impl ChanID {
    fn new_copy(src: MemID, dst: MemID) -> Self {
        ChanID {
            src: Some(src),
            dst: Some(dst),
        }
    }
    fn new_fill(dst: MemID) -> Self {
        ChanID {
            src: None,
            dst: Some(dst),
        }
    }
    fn new_deppart() -> Self {<commit_msg>legion_prof_rs: Match new Legion Prof sorting behavior.<commit_after>use std::cmp::{max, Reverse};
use std::fmt;
use std::collections::{BTreeMap, BTreeSet, BinaryHeap};
use std::convert::TryFrom;

use derive_more::{Add, From, LowerHex, Sub};
use num_enum::TryFromPrimitive;

use crate::serialize::Record;

const TASK_GRANULARITY_THRESHOLD: u64 = 10 * 1000;


// Make sure this is up to date with lowlevel.h
#[derive(Debug, Copy, Clone, Eq, PartialEq, Ord, PartialOrd, TryFromPrimitive)]
#[repr(i32)]
pub enum ProcKind {
    GPU = 1,
    CPU = 2,
    Utility = 3,
    IO = 4,
    ProcGroup = 5,
    OpenMP = 6,
    Python = 7,
}

// Make sure this is up to date with lowlevel.h
#[derive(Debug, Copy, Clone, Eq, PartialEq, TryFromPrimitive)]
#[repr(i32)]
pub enum MemKind {
    NoMemKind = 0,
    Global = 1,
    System = 2,
    Registered = 3,
    Socket = 4,
    ZeroCopy = 5,
    Framebuffer = 6,
    Disk = 7,
    HDF5 = 8,
    File = 9,
    L3Cache = 10,
    L2Cache = 11,
    L1Cache = 12,
}

// Make sure this is up to date with lowlevel.h
#[derive(Debug, Copy, Clone, Eq, PartialEq, TryFromPrimitive)]
#[repr(i32)]
pub enum DepPartKind {
    Union = 0,
    Unions = 1,
    UnionReduction = 2,
    Intersection = 3,
    Intersections = 4,
    IntersectionReduction = 5,
    Difference = 6,
    Differences = 7,
    EqualPartition = 8,
    PartitionByField = 9,
    PartitionByImage = 10,
    PartitionByImageRange = 11,
    PartitionByPreimage = 12,
    PartitionByPreimageRange = 13,
    CreateAssociation = 14,
    PartitionByWeights = 15,
}

// Make sure this is up to date with lowlevel.h
#[derive(Debug, Copy, Clone, Eq, PartialEq, TryFromPrimitive)]
#[repr(u32)]
pub enum DimKind {
    DimX = 0,
    DimY = 1,
    DimZ = 2,
    DimW = 3,
    DimV = 4,
    DimU = 5,
    DimT = 6,
    DimS = 7,
    DimR = 8,
    DimF = 9,
    InnerDimX = 10,
    OuterDimX = 11,
    InnerDimY = 12,
    OuterDimY = 13,
    InnerDimZ = 14,
    OuterDimZ = 15,
    InnerDimW = 16,
    OuterDimW = 17,
    InnerDimV = 18,
    OuterDimV = 19,
    InnerDimU = 20,
    OuterDimU = 21,
    InnerDimT = 22,
    OuterDimT = 23,
    InnerDimS = 24,
    OuterDimS = 25,
    InnerDimR = 26,
    OuterDimR = 27,
}

#[derive(Debug, Copy, Clone, PartialEq, Eq, PartialOrd, Ord, Default, Add, Sub, From)]
pub struct Timestamp(pub u64 /* ns */);

impl fmt::Display for Timestamp {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        // Time is stored in nanoseconds. But it is displayed in microseconds.
        let nanoseconds = self.0;
        let divisor = 1000;
        let microseconds = nanoseconds / divisor;
        let remainder = nanoseconds % divisor;
        write!(f, "{}.{:0>3}", microseconds, remainder)
    }
}

#[derive(Debug, Copy, Clone)]
pub enum ProcEntry {
    Task(OpID),
    MetaTask(OpID, VariantID, usize),
    MapperCall(usize),
    RuntimeCall(usize),
    ProfTask(usize)
}

#[derive(Debug, Copy, Clone)]
pub struct ProcPoint {
    pub time: Timestamp,
    // Secondary sort_key, used for breaking ties in sorting
    // In practice, we plan for this to be a nanosecond timestamp,
    // like the time field above.
    pub secondary_sort_key: u64,
    pub entry: ProcEntry,
    pub first: bool,
}

impl ProcPoint {
    fn new(time: Timestamp, entry: ProcEntry, first: bool, secondary_sort_key: u64) -> Self {
        ProcPoint { time, entry, first , secondary_sort_key: secondary_sort_key}
    }
    pub fn time_key(&self) -> (u64, u8, u64) {
        (self.time.0, if self.first { 0 } else { 1 }, self.secondary_sort_key)
    }
}

#[derive(Debug, Copy, Clone, PartialEq, Eq, PartialOrd, Ord, LowerHex)]
pub struct ProcID(pub u64);

#[derive(Debug, Copy, Clone, PartialEq, Eq, PartialOrd, Ord)]
pub struct NodeID(pub u64);

#[derive(Debug, Copy, Clone, PartialEq, Eq, PartialOrd, Ord)]
pub struct EventID(pub u64);

#[derive(Debug)]
pub struct Proc {
    pub proc_id: ProcID,
    pub kind: ProcKind,
    pub node_id: NodeID,
    pub proc_in_node: u64,
    app_ranges: Vec<()>,
    last_time: Option<Timestamp>,
    pub tasks: BTreeMap<OpID, Task>,
    pub meta_tasks: BTreeMap<(OpID, VariantID), Vec<MetaTask>>,
    pub mapper_calls: Vec<MapperCall>,
    pub runtime_calls: Vec<RuntimeCall>,
    pub prof_tasks: Vec<ProfTask>,
    pub max_levels: u32,
    pub max_levels_ready: u32,
    pub time_points: Vec<ProcPoint>,
    pub util_time_points: Vec<ProcPoint>,
}

impl Proc {
    fn new(proc_id: ProcID, kind: ProcKind) -> Self {
        Proc {
            proc_id,
            kind,
            // PROCESSOR:   tag:8 = 0x1d, owner_node:16,   (unused):28, proc_idx: 12
            // owner_node = proc_id[55:40]
            // proc_idx = proc_id[11:0]
            node_id: NodeID((proc_id.0 >> 40) & ((1 << 16) - 1)),
            proc_in_node: (proc_id.0) & ((1 << 12) - 1),
            app_ranges: Vec::new(),
            last_time: None,
            tasks: BTreeMap::new(),
            meta_tasks: BTreeMap::new(),
            mapper_calls: Vec::new(),
            runtime_calls: Vec::new(),
            prof_tasks: Vec::new(),
            max_levels: 0,
            max_levels_ready: 0,
            time_points: Vec::new(),
            util_time_points: Vec::new(),
        }
    }

    pub fn is_empty(&self) -> bool {
        self.tasks.is_empty() && self.mapper_calls.is_empty() && self.runtime_calls.is_empty() && self.prof_tasks.is_empty()
    }

    pub fn entry(&self, entry: ProcEntry) -> (&Base, &TimeRange, &Waiters) {
        match entry {
            ProcEntry::Task(op_id) => {
                let task = self.tasks.get(&op_id).unwrap();
                (&task.base, &task.time_range, &task.waiters)
            }
            ProcEntry::MetaTask(op_id, variant_id, idx) => {
                let task = &self.meta_tasks.get(&(op_id, variant_id)).unwrap()[idx];
                (&task.base, &task.time_range, &task.waiters)
            }
            ProcEntry::MapperCall(idx) => {
                let call = &self.mapper_calls[idx];
                (&call.base, &call.time_range, &call.waiters)
            }
            ProcEntry::RuntimeCall(idx) => {
                let call = &self.runtime_calls[idx];
                (&call.base, &call.time_range, &call.waiters)
            }
            ProcEntry::ProfTask(idx) => {
                let task = &self.prof_tasks.get(idx).unwrap();
                (&task.base, &task.time_range, &task.waiters)
            }
        }
    }

    pub fn entry_mut(&mut self, entry: ProcEntry) -> (&mut Base, &mut TimeRange, &mut Waiters) {
        match entry {
            ProcEntry::Task(op_id) => {
                let task = self.tasks.get_mut(&op_id).unwrap();
                (&mut task.base, &mut task.time_range, &mut task.waiters)
            }
            ProcEntry::MetaTask(op_id, variant_id, idx) => {
                let task = &mut self.meta_tasks.get_mut(&(op_id, variant_id)).unwrap()[idx];
                (&mut task.base, &mut task.time_range, &mut task.waiters)
            }
            ProcEntry::MapperCall(idx) => {
                let call = &mut self.mapper_calls[idx];
                (&mut call.base, &mut call.time_range, &mut call.waiters)
            }
            ProcEntry::RuntimeCall(idx) => {
                let call = &mut self.runtime_calls[idx];
                (&mut call.base, &mut call.time_range, &mut call.waiters)
            }
            ProcEntry::ProfTask(idx) => {
                let task = self.prof_tasks.get_mut(idx).unwrap();
                (&mut task.base, &mut task.time_range, &mut task.waiters)
            }
        }
    }

    pub fn entry_base(&self, entry: ProcEntry) -> &Base {
        self.entry(entry).0
    }

    pub fn entry_base_mut(&mut self, entry: ProcEntry) -> &mut Base {
        self.entry_mut(entry).0
    }

    fn trim_time_range(&mut self, start: Timestamp, stop: Timestamp) {
        let mut removed_tasks = Vec::new();
        for (op_id, task) in self.tasks.iter_mut() {
            task.trim_time_range(start, stop);
            if task.time_range.was_removed {
                removed_tasks.push(*op_id);
            }
        }
        for op_id in removed_tasks {
            self.tasks.remove(&op_id);
        }
        for mapper_call in &mut self.mapper_calls {
            mapper_call.trim_time_range(start, stop);
        }
        self.mapper_calls.retain(|t| !t.time_range.was_removed);
        for runtime_call in &mut self.runtime_calls {
            runtime_call.trim_time_range(start, stop);
        }
        self.runtime_calls.retain(|t| !t.time_range.was_removed);
    }

    fn sort_time_range(&mut self, last_time: Timestamp) {
        fn add(
            time: &TimeRange,
            entry: ProcEntry,
            all_points: &mut Vec<ProcPoint>,
            points: &mut Vec<ProcPoint>,
            util_points: &mut Vec<ProcPoint>,
        ) {
            let start = time.start.unwrap();
            let stop = time.stop.unwrap();
            let ready = time.ready;
            if stop - start > TASK_GRANULARITY_THRESHOLD.into() && !ready.is_none() {
                all_points.push(ProcPoint::new(ready.unwrap(), entry, true, start.0));
                all_points.push(ProcPoint::new(stop, entry, false, 0));
            } else {
                all_points.push(ProcPoint::new(start, entry, true, 0));
                all_points.push(ProcPoint::new(stop, entry, false, 0));
            }

            points.push(ProcPoint::new(start, entry, true, 0));
            points.push(ProcPoint::new(stop, entry, false, 0));

            util_points.push(ProcPoint::new(start, entry, true, 0));
            util_points.push(ProcPoint::new(stop, entry, false, 0));
        }
        fn add_waiters(waiters: &Waiters, entry: ProcEntry, util_points: &mut Vec<ProcPoint>) {
            for wait in &waiters.wait_intervals {
                util_points.push(ProcPoint::new(wait.start, entry, false, 0));
                util_points.push(ProcPoint::new(wait.end, entry, true, 0));
            }
        }

        let mut all_points = Vec::new();
        let mut points = Vec::new();
        let mut util_points = Vec::new();

        for task in self.tasks.values() {
            let time = &task.time_range;
            let entry = ProcEntry::Task(task.op_id);
            add(&time, entry, &mut all_points, &mut points, &mut util_points);
            add_waiters(&task.waiters, entry, &mut util_points);
        }
        for tasks in self.meta_tasks.values() {
            for (idx, task) in tasks.iter().enumerate() {
                let time = &task.time_range;
                let entry = ProcEntry::MetaTask(task.op_id, task.variant_id, idx);
                add(&time, entry, &mut all_points, &mut points, &mut util_points);
                add_waiters(&task.waiters, entry, &mut util_points);
            }
        }
        for (idx, mapper_call) in self.mapper_calls.iter().enumerate() {
            let time = &mapper_call.time_range;
            let entry = ProcEntry::MapperCall(idx);
            add(&time, entry, &mut all_points, &mut points, &mut util_points);
            add_waiters(&mapper_call.waiters, entry, &mut util_points);
        }
        for (idx, runtime_call) in self.runtime_calls.iter().enumerate() {
            let time = &runtime_call.time_range;
            let entry = ProcEntry::RuntimeCall(idx);
            add(&time, entry, &mut all_points, &mut points, &mut util_points);
            add_waiters(&runtime_call.waiters, entry, &mut util_points);
        }

        for (idx, prof_task) in self.prof_tasks.iter().enumerate() {
            let time = &prof_task.time_range;
            let entry = ProcEntry::ProfTask(idx);
            add(&time, entry, &mut all_points, &mut points, &mut util_points);
            add_waiters(&prof_task.waiters, entry, &mut util_points);
        }

        points.sort_by(|a, b| a.time_key().cmp(&b.time_key()));
        util_points.sort_by(|a, b| a.time_key().cmp(&b.time_key()));

        // Hack: This is a max heap so reverse the values as they go in.
        let mut free_levels = BinaryHeap::<Reverse<u32>>::new();
        for point in &points {
            if point.first {
                let level = if let Some(level) = free_levels.pop() {
                    level.0
                } else {
                    self.max_levels += 1;
                    self.max_levels
                };
                self.entry_base_mut(point.entry).set_level(level);
            } else {
                let level = self.entry_base(point.entry).level.unwrap();
                free_levels.push(Reverse(level));
            }
        }

        all_points.sort_by(|a, b| a.time_key().cmp(&b.time_key()));

        // Hack: This is a max heap so reverse the values as they go in.
        let mut free_levels_ready = BinaryHeap::<Reverse<u32>>::new();
        for point in &all_points {
            if point.first {
                let level = if let Some(level) = free_levels_ready.pop() {
                    level.0
                } else {
                    self.max_levels_ready += 1;
                    self.max_levels_ready
                };
                self.entry_base_mut(point.entry).set_level_ready(level);
            } else {
                let level = self.entry_base(point.entry).level_ready.unwrap();
                free_levels_ready.push(Reverse(level));
            }
        }

        self.time_points = points;
        self.util_time_points = util_points;
    }
}

#[derive(Debug, Copy, Clone, PartialEq, Eq, PartialOrd, Ord)]
pub struct MemID(pub u64);

impl MemID {
    pub fn node_id(&self) -> NodeID {
        NodeID((self.0 >> 40) & ((1 << 16) - 1))
    }
}

#[derive(Debug)]
pub struct Mem {
    mem_id: MemID,
    kind: MemKind,
    capacity: u64,
    node_id: NodeID,
    mem_in_node: u64,
    instances: BTreeSet<()>,
    time_points: Vec<()>,
    max_live_instances: Option<u64>,
    last_time: Option<Timestamp>,
    affinity: Option<u64>,
}

impl Mem {
    fn new(mem_id: MemID, kind: MemKind, capacity: u64) -> Self {
        Mem {
            mem_id,
            kind,
            capacity,
            node_id: NodeID((mem_id.0 >> 40) & ((1 << 16) - 1)),
            mem_in_node: (mem_id.0) & ((1 << 12) - 1),
            instances: BTreeSet::new(),
            time_points: Vec::new(),
            max_live_instances: None,
            last_time: None,
            affinity: None,
        }
    }
    fn trim_time_range(&mut self, start: Timestamp, stop: Timestamp) {}
    fn sort_time_range(&mut self, last_time: Timestamp) {}
}

#[derive(Debug)]
pub struct MemProcAffinity {
    mem_id: MemID,
    proc_ids: Vec<ProcID>,
}

impl MemProcAffinity {
    fn new(mem_id: MemID) -> Self {
        MemProcAffinity {
            mem_id,
            proc_ids: Vec::new(),
        }
    }
    fn add_proc_id(&mut self, proc_id: ProcID) {
        self.proc_ids.push(proc_id);
    }
}

#[derive(Debug, Copy, Clone)]
pub enum ChanEntry {
    Copy(usize),
    Fill(usize),
    DepPart(usize),
}

#[derive(Debug, Copy, Clone)]
pub struct ChanPoint {
    pub time: Timestamp,
    pub entry: ChanEntry,
    pub first: bool,
}

impl ChanPoint {
    fn new(time: Timestamp, entry: ChanEntry, first: bool) -> Self {
        ChanPoint { time, entry, first }
    }
    pub fn time_key(&self) -> (u64, u8) {
        (self.time.0, if self.first { 0 } else { 1 })
    }
}

#[derive(Debug, Copy, Clone, PartialEq, Eq, PartialOrd, Ord)]
pub struct ChanID {
    pub src: Option<MemID>,
    pub dst: Option<MemID>,
}

impl ChanID {
    fn new_copy(src: MemID, dst: MemID) -> Self {
        ChanID {
            src: Some(src),
            dst: Some(dst),
        }
    }
    fn new_fill(dst: MemID) -> Self {
        ChanID {
            src: None,
            dst: Some(dst),
        }
    }
    fn new_deppart() -> Self {
{'fraction_non_alphanumeric': 0.1310300703774792, 'fraction_numerical': 0.014395393474088292, 'mean_word_length': 2.100982045432001, 'pattern_counts': {'":': 0, '<': 77, '<?xml version=': 0, '>': 127, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '908799', 'n_tokens_mistral': 10774, 'n_tokens_neox': 10268, 'n_words': 2397}