From fc40421dc973fac623133a219e092bb67ef8220a Mon Sep 17 00:00:00 2001
From: Rahul Kumar <rahulkumar@berkeley.edu>
Date: Wed, 12 Jun 2024 12:11:29 -0700
Subject: [PATCH] feat(schematics): support SCIR netlist exports with multiple
 top cells (#424)

---
 substrate/src/context.rs        | 12 ++++++++-
 substrate/src/schematic/conv.rs | 23 +++++++++++++++++
 substrate/src/schematic/mod.rs  |  6 +++++
 tests/src/schematic.rs          | 46 +++++++++++++++++++++++++++++++++
 4 files changed, 86 insertions(+), 1 deletion(-)

diff --git a/substrate/src/context.rs b/substrate/src/context.rs
index a052b5c1..737bbf35 100644
--- a/substrate/src/context.rs
+++ b/substrate/src/context.rs
@@ -35,7 +35,7 @@ use crate::pdk::layers::LayerId;
 use crate::pdk::layers::Layers;
 use crate::pdk::layers::{GdsLayerSpec, InstalledLayers};
 use crate::pdk::Pdk;
-use crate::schematic::conv::{ConvError, RawLib};
+use crate::schematic::conv::{export_multi_top_scir_lib, ConvError, RawLib};
 use crate::schematic::schema::{FromSchema, Schema};
 use crate::schematic::{
     Cell as SchematicCell, CellCacheKey, CellHandle as SchematicCellHandle, CellId, CellMetadata,
@@ -413,6 +413,16 @@ impl Context {
         raw.to_scir_lib()
     }
 
+    /// Export the given cells and all their subcells as a SCIR library.
+    ///
+    /// Returns a SCIR library and metadata for converting between SCIR and Substrate formats.
+    pub fn export_scir_all<S: Schema + ?Sized>(
+        &self,
+        cells: &[&crate::schematic::RawCell<S>],
+    ) -> Result<RawLib<S>, ConvError> {
+        export_multi_top_scir_lib(cells)
+    }
+
     /// Simulate the given testbench.
     ///
     /// The simulator must be installed in the context.
diff --git a/substrate/src/schematic/conv.rs b/substrate/src/schematic/conv.rs
index 4522e9c1..06b8ba15 100644
--- a/substrate/src/schematic/conv.rs
+++ b/substrate/src/schematic/conv.rs
@@ -492,6 +492,9 @@ impl<S: Schema + ?Sized> RawCell<S> {
     /// Export this cell and all subcells as a SCIR library.
     ///
     /// Returns the SCIR library and metadata for converting between SCIR and Substrate formats.
+    ///
+    /// Consider using [`export_multi_top_scir_lib`] if you need to export multiple cells
+    /// to the same SCIR library.
     pub(crate) fn to_scir_lib(&self) -> Result<RawLib<S>, ConvError> {
         let mut lib_ctx = ScirLibExportContext::new();
         let scir_id = self.to_scir_cell(&mut lib_ctx)?;
@@ -754,3 +757,23 @@ pub enum ConvError {
     #[error("unsupported primitive")]
     UnsupportedPrimitive,
 }
+
+/// Export a collection of cells and all their subcells as a SCIR library.
+///
+/// Returns the SCIR library and metadata for converting between SCIR and Substrate formats.
+/// The resulting SCIR library will **not** have a top cell set.
+/// If you want a SCIR library with a known top cell, consider using [`RawCell::to_scir_lib`] instead.
+pub(crate) fn export_multi_top_scir_lib<S: Schema + ?Sized>(
+    cells: &[&RawCell<S>],
+) -> Result<RawLib<S>, ConvError> {
+    let mut lib_ctx = ScirLibExportContext::new();
+
+    for &cell in cells {
+        cell.to_scir_cell(&mut lib_ctx)?;
+    }
+
+    Ok(RawLib {
+        scir: lib_ctx.lib.build()?,
+        conv: lib_ctx.conv.build(),
+    })
+}
diff --git a/substrate/src/schematic/mod.rs b/substrate/src/schematic/mod.rs
index 1f2a5316..24a08463 100644
--- a/substrate/src/schematic/mod.rs
+++ b/substrate/src/schematic/mod.rs
@@ -619,6 +619,12 @@ impl<S: Schema, B: ExportsNestedData> SchemaCellHandle<S, B> {
     pub fn cell(&self) -> &Cell<B> {
         self.cell.cell()
     }
+
+    /// Returns the raw cell.
+    pub fn raw(&self) -> Arc<RawCell<S>> {
+        let val = self.handle.unwrap_inner();
+        val.raw.clone()
+    }
 }
 
 impl<S: Schema + ?Sized, B: ExportsNestedData> Deref for SchemaCellHandle<S, B> {
diff --git a/tests/src/schematic.rs b/tests/src/schematic.rs
index ef912e06..ceb72d27 100644
--- a/tests/src/schematic.rs
+++ b/tests/src/schematic.rs
@@ -52,6 +52,52 @@ fn can_generate_vdivider_schematic() {
     assert_eq!(vdiv.instances().count(), 2);
 }
 
+#[test]
+fn can_generate_multi_top_scir() {
+    let ctx = Context::new();
+    let vdivider1 = Vdivider {
+        r1: Resistor::new(300),
+        r2: Resistor::new(100),
+    };
+    let vdivider2 = Vdivider {
+        r1: Resistor::new(500),
+        r2: Resistor::new(600),
+    };
+    let vdiv1 = ctx.generate_schematic::<Spice, _>(vdivider1);
+    let vdiv2 = ctx.generate_schematic::<Spice, _>(vdivider2);
+    let RawLib { scir, conv: _ } = ctx.export_scir_all(&[&vdiv1.raw(), &vdiv2.raw()]).unwrap();
+    assert_eq!(scir.cells().count(), 2);
+    let issues = scir.validate();
+    println!("Library:\n{:#?}", scir);
+    println!("Issues = {:#?}", issues);
+    assert_eq!(issues.num_errors(), 0);
+    assert_eq!(issues.num_warnings(), 0);
+
+    let vdiv = scir.cell_named("vdivider_300_100");
+    let port_names: HashSet<ArcStr> = vdiv
+        .ports()
+        .map(|p| vdiv.signal(p.signal()).name.clone())
+        .collect();
+    assert_eq!(port_names.len(), 3);
+    assert!(port_names.contains("pwr_vdd"));
+    assert!(port_names.contains("pwr_vss"));
+    assert!(port_names.contains("out"));
+    assert_eq!(vdiv.ports().count(), 3);
+    assert_eq!(vdiv.instances().count(), 2);
+
+    let vdiv = scir.cell_named("vdivider_500_600");
+    let port_names: HashSet<ArcStr> = vdiv
+        .ports()
+        .map(|p| vdiv.signal(p.signal()).name.clone())
+        .collect();
+    assert_eq!(port_names.len(), 3);
+    assert!(port_names.contains("pwr_vdd"));
+    assert!(port_names.contains("pwr_vss"));
+    assert!(port_names.contains("out"));
+    assert_eq!(vdiv.ports().count(), 3);
+    assert_eq!(vdiv.instances().count(), 2);
+}
+
 #[test]
 fn can_generate_flattened_vdivider_schematic() {
     let ctx = Context::new();