Add binary coverage merging

src/agent/coverage/Cargo.toml

@@ -24,3 +24,4 @@ procfs = { version = "0.12", default-features = false, features=["flate2"] }
 
 [dev-dependencies]
 clap = { version = "4.0", features = ["derive"] }
+pretty_assertions = "1.3.0"

src/agent/coverage/src/binary.rs

@@ -3,7 +3,7 @@
 
 use std::collections::{BTreeMap, BTreeSet};
 
-use anyhow::{bail, Result};
+use anyhow::Result;
 use debuggable_module::Module;
 pub use debuggable_module::{block, path::FilePath, Offset};
 use symbolic::debuginfo::Object;
@@ -16,24 +16,62 @@ pub struct BinaryCoverage {
     pub modules: BTreeMap<FilePath, ModuleBinaryCoverage>,
 }
 
+impl BinaryCoverage {
+    pub fn add(&mut self, rhs: &Self) {
+        for (path, rhs_module) in &rhs.modules {
+            let module = self.modules.entry(path.clone()).or_default();
+            module.add(rhs_module);
+        }
+    }
+
+    pub fn merge(&mut self, rhs: &Self) {
+        for (path, rhs_module) in &rhs.modules {
+            let module = self.modules.entry(path.clone()).or_default();
+            module.merge(rhs_module);
+        }
+    }
+}
+
 #[derive(Clone, Debug, Default, Eq, PartialEq)]
 pub struct ModuleBinaryCoverage {
     pub offsets: BTreeMap<Offset, Count>,
 }
 
 impl ModuleBinaryCoverage {
-    pub fn increment(&mut self, offset: Offset) -> Result<()> {
-        if let Some(count) = self.offsets.get_mut(&offset) {
-            count.increment();
-        } else {
-            bail!("unknown coverage offset: {offset:x}");
-        };
-
-        Ok(())
+    pub fn increment(&mut self, offset: Offset) {
+        let count = self.offsets.entry(offset).or_default();
+        count.increment();
+    }
+
+    pub fn add(&mut self, rhs: &Self) {
+        for (&offset, &rhs_count) in &rhs.offsets {
+            let count = self.offsets.entry(offset).or_default();
+            *count += rhs_count;
+        }
+    }
+
+    pub fn merge(&mut self, rhs: &Self) {
+        for (&offset, &rhs_count) in &rhs.offsets {
+            let count = self.offsets.entry(offset).or_default();
+            *count = Count::max(*count, rhs_count)
+        }
     }
 }
 
-#[derive(Clone, Copy, Debug, Eq, PartialEq)]
+impl<O> From<O> for ModuleBinaryCoverage
+where
+    O: IntoIterator<Item = Offset>,
+{
+    fn from(offsets: O) -> Self {
+        let offsets = offsets.into_iter().map(|o| (o, Count(0)));
+
+        let mut coverage = Self::default();
+        coverage.offsets.extend(offsets);
+        coverage
+    }
+}
+
+#[derive(Clone, Copy, Debug, Default, Eq, PartialEq)]
 pub struct Count(pub u32);
 
 impl Count {
@@ -44,6 +82,24 @@ impl Count {
     pub fn reached(&self) -> bool {
         self.0 > 0
     }
+
+    pub fn max(self, rhs: Self) -> Self {
+        Count(u32::max(self.0, rhs.0))
+    }
+}
+
+impl std::ops::Add for Count {
+    type Output = Self;
+
+    fn add(self, rhs: Self) -> Self {
+        Count(self.0.saturating_add(rhs.0))
+    }
+}
+
+impl std::ops::AddAssign for Count {
+    fn add_assign(&mut self, rhs: Self) {
+        *self = *self + rhs;
+    }
 }
 
 pub fn find_coverage_sites<'data>(
@@ -81,10 +137,7 @@ pub fn find_coverage_sites<'data>(
         }
     }
 
-    let mut coverage = ModuleBinaryCoverage::default();
-    coverage
-        .offsets
-        .extend(offsets.into_iter().map(|o| (o, Count(0))));
+    let coverage = ModuleBinaryCoverage::from(offsets.into_iter());
 
     Ok(coverage)
 }
@@ -94,3 +147,6 @@ impl AsRef<BTreeMap<Offset, Count>> for ModuleBinaryCoverage {
         &self.offsets
     }
 }
+
+#[cfg(test)]
+mod tests;

src/agent/coverage/src/binary/tests.rs

@@ -0,0 +1,171 @@
+// Copyright (c) Microsoft Corporation.
+// Licensed under the MIT License.
+
+use anyhow::Result;
+use debuggable_module::path::FilePath;
+use pretty_assertions::assert_eq;
+
+use crate::binary::{Count, Offset};
+
+use super::*;
+
+macro_rules! module {
+    ( $( $offset: expr => $count: expr, )* ) => {{
+        let mut module = ModuleBinaryCoverage::default();
+
+        $(
+            module.offsets.insert(Offset($offset), Count($count));
+        )*
+
+        module
+    }}
+}
+
+macro_rules! coverage {
+    ( $( $path: expr => { $( $offset: expr => $count: expr, )* }, )* ) => {{
+        let mut coverage = BinaryCoverage::default();
+
+        $(
+            let path = FilePath::new($path)?;
+            let module = module! { $( $offset => $count, )* };
+            coverage.modules.insert(path, module);
+        )*
+
+        coverage
+    }}
+}
+
+#[test]
+fn test_module_increment() -> Result<()> {
+    let mut module = module! {
+        1 => 1,
+        2 => 0,
+    };
+
+    module.increment(Offset(2));
+
+    assert_eq!(
+        module,
+        module! {
+            1 => 1,
+            2 => 1,
+        }
+    );
+
+    module.increment(Offset(2));
+
+    assert_eq!(
+        module,
+        module! {
+            1 => 1,
+            2 => 2,
+        }
+    );
+
+    module.increment(Offset(3));
+
+    assert_eq!(
+        module,
+        module! {
+            1 => 1,
+            2 => 2,
+            3 => 1,
+        }
+    );
+
+    Ok(())
+}
+
+#[test]
+fn test_coverage_add() -> Result<()> {
+    let mut coverage = coverage! {
+        "main.exe" => {
+            1 => 1,
+            2 => 0,
+            3 => 1,
+            4 => 0,
+        },
+        "old.dll" => {
+            1 => 0,
+        },
+    };
+
+    coverage.add(&coverage! {
+        "main.exe" => {
+            1 => 1,
+            2 => 1,
+            5 => 1,
+        },
+        "new.dll" => {
+            1 => 1,
+        },
+    });
+
+    assert_eq!(
+        coverage,
+        coverage! {
+            "main.exe" => {
+                1 => 2,
+                2 => 1,
+                3 => 1,
+                4 => 0,
+                5 => 1,
+            },
+            "old.dll" => {
+                1 => 0,
+            },
+            "new.dll" => {
+                1 => 1,
+            },
+        }
+    );
+
+    Ok(())
+}
+
+#[test]
+fn test_coverage_merge() -> Result<()> {
+    let mut coverage = coverage! {
+        "main.exe" => {
+            1 => 1,
+            2 => 0,
+            3 => 1,
+            4 => 0,
+        },
+        "old.dll" => {
+            1 => 0,
+        },
+    };
+
+    coverage.merge(&coverage! {
+        "main.exe" => {
+            1 => 1,
+            2 => 1,
+            5 => 1,
+        },
+        "new.dll" => {
+            1 => 1,
+        },
+    });
+
+    assert_eq!(
+        coverage,
+        coverage! {
+            "main.exe" => {
+                1 => 1,
+                2 => 1,
+                3 => 1,
+                4 => 0,
+                5 => 1,
+            },
+            "old.dll" => {
+                1 => 0,
+            },
+            "new.dll" => {
+                1 => 1,
+            },
+        }
+    );
+
+    Ok(())
+}

src/agent/coverage/src/record/linux.rs

@@ -48,7 +48,7 @@ impl<'data> LinuxRecorder<'data> {
         if let Some(image) = context.find_image_for_addr(addr) {
             if let Some(coverage) = self.coverage.modules.get_mut(image.path()) {
                 let offset = addr.offset_from(image.base())?;
-                coverage.increment(offset)?;
+                coverage.increment(offset);
             } else {
                 bail!("coverage not initialized for module {}", image.path());
             }

src/agent/coverage/src/record/windows.rs

@@ -75,7 +75,7 @@ impl<'data> WindowsRecorder<'data> {
             .get_mut(&breakpoint.module)
             .ok_or_else(|| anyhow!("coverage not initialized for module: {}", breakpoint.module))?;
 
-        coverage.increment(breakpoint.offset)?;
+        coverage.increment(breakpoint.offset);
 
         Ok(())
     }