Merge pull request #580 from EnergySystemsModellingLab/id-tweaks

Various tweaks to ID code

Author: alexdewar

src/id.rs

@@ -1,17 +1,14 @@
 //! Code for handling IDs
 use anyhow::{Context, Result};
-use indexmap::IndexSet;
+use indexmap::{IndexMap, IndexSet};
+use std::borrow::Borrow;
 use std::collections::HashSet;
+use std::fmt::Display;
+use std::hash::Hash;
 
 /// A trait alias for ID types
-pub trait IDLike:
-    Eq + std::hash::Hash + std::borrow::Borrow<str> + Clone + std::fmt::Display + From<String>
-{
-}
-impl<T> IDLike for T where
-    T: Eq + std::hash::Hash + std::borrow::Borrow<str> + Clone + std::fmt::Display + From<String>
-{
-}
+pub trait IDLike: Eq + Hash + Borrow<str> + Clone + Display + From<String> {}
+impl<T> IDLike for T where T: Eq + Hash + Borrow<str> + Clone + Display + From<String> {}
 
 macro_rules! define_id_type {
     ($name:ident) => {
@@ -78,43 +75,25 @@ pub(crate) use define_id_getter;
 
 /// A data structure containing a set of IDs
 pub trait IDCollection<ID: IDLike> {
-    /// Get the ID from the collection by its string representation.
-    ///
-    /// # Arguments
-    ///
-    /// * `id` - The string representation of the ID
-    ///
-    /// # Returns
-    ///
-    /// A copy of the ID in `self`, or an error if not found.
-    fn get_id_by_str(&self, id: &str) -> Result<ID>;
-
     /// Check if the ID is in the collection, returning a copy of it if found.
     ///
     /// # Arguments
     ///
-    /// * `id` - The ID to check
+    /// * `id` - The ID to check (can be string or ID type)
     ///
     /// # Returns
     ///
     /// A copy of the ID in `self`, or an error if not found.
-    fn get_id(&self, id: &ID) -> Result<ID>;
+    fn get_id<T: Borrow<str> + Display + ?Sized>(&self, id: &T) -> Result<&ID>;
 }
 
 macro_rules! define_id_methods {
     () => {
-        fn get_id_by_str(&self, id: &str) -> Result<ID> {
-            let found = self
-                .get(id)
-                .with_context(|| format!("Unknown ID {id} found"))?;
-            Ok(found.clone())
-        }
-
-        fn get_id(&self, id: &ID) -> Result<ID> {
+        fn get_id<T: Borrow<str> + Display + ?Sized>(&self, id: &T) -> Result<&ID> {
             let found = self
                 .get(id.borrow())
                 .with_context(|| format!("Unknown ID {id} found"))?;
-            Ok(found.clone())
+            Ok(found)
         }
     };
 }
@@ -126,3 +105,12 @@ impl<ID: IDLike> IDCollection<ID> for HashSet<ID> {
 impl<ID: IDLike> IDCollection<ID> for IndexSet<ID> {
     define_id_methods!();
 }
+
+impl<ID: IDLike, V> IDCollection<ID> for IndexMap<ID, V> {
+    fn get_id<T: Borrow<str> + Display + ?Sized>(&self, id: &T) -> Result<&ID> {
+        let (found, _) = self
+            .get_key_value(id.borrow())
+            .with_context(|| format!("Unknown ID {id} found"))?;
+        Ok(found)
+    }
+}

src/input/agent/commodity_portion.rs

@@ -2,6 +2,7 @@
 use super::super::*;
 use crate::agent::{AgentCommodityPortionsMap, AgentID, AgentMap};
 use crate::commodity::{CommodityID, CommodityMap, CommodityType};
+use crate::id::IDCollection;
 use crate::region::RegionID;
 use crate::year::parse_year_str;
 use anyhow::{ensure, Context, Result};
@@ -66,9 +67,7 @@ where
     for agent_commodity_portion_raw in iter {
         // Get agent ID
         let agent_id_raw = agent_commodity_portion_raw.agent_id.as_str();
-        let (id, _agent) = agents
-            .get_key_value(agent_id_raw)
-            .with_context(|| format!("Invalid agent ID {agent_id_raw}"))?;
+        let id = agents.get_id(agent_id_raw)?;
 
         // Get/create entry for agent
         let entry = agent_commodity_portions
@@ -77,9 +76,7 @@ where
 
         // Insert portion for the commodity/year(s)
         let commodity_id_raw = agent_commodity_portion_raw.commodity_id.as_str();
-        let (commodity_id, _commodity) = commodities
-            .get_key_value(commodity_id_raw)
-            .with_context(|| format!("Invalid commodity ID {commodity_id_raw}"))?;
+        let commodity_id = commodities.get_id(commodity_id_raw)?;
         let years = parse_year_str(&agent_commodity_portion_raw.years, milestone_years)?;
         for year in years {
             try_insert(

src/input/agent/cost_limit.rs

@@ -61,7 +61,7 @@ where
         let years = parse_year_str(&agent_cost_limits_raw.years, milestone_years)?;
 
         // Get agent ID
-        let agent_id = agent_ids.get_id_by_str(&agent_cost_limits_raw.agent_id)?;
+        let agent_id = agent_ids.get_id(&agent_cost_limits_raw.agent_id)?;
 
         // Get or create entry in the map
         let entry = map.entry(agent_id.clone()).or_default();

src/input/agent/search_space.rs

@@ -52,18 +52,16 @@ impl AgentSearchSpaceRaw {
         let search_space = Rc::new(parse_search_space_str(&self.search_space, process_ids)?);
 
         // Get commodity
-        let commodity_id = commodity_ids.get_id_by_str(&self.commodity_id)?;
+        let commodity_id = commodity_ids.get_id(&self.commodity_id)?;
 
         // Check that the year is a valid milestone year
         let year = parse_year_str(&self.years, milestone_years)?;
 
-        let (agent_id, _) = agents
-            .get_key_value(self.agent_id.as_str())
-            .context("Invalid agent ID")?;
+        let agent_id = agents.get_id(&self.agent_id)?;
 
         Ok(AgentSearchSpace {
             agent_id: agent_id.clone(),
-            commodity_id,
+            commodity_id: commodity_id.clone(),
             years: year,
             search_space,
         })
@@ -86,7 +84,7 @@ fn parse_search_space_str(
     } else {
         search_space
             .split(';')
-            .map(|id| process_ids.get_id_by_str(id.trim()))
+            .map(|id| Ok(process_ids.get_id(id.trim())?.clone()))
             .try_collect()
     }
 }

src/input/asset.rs

@@ -69,16 +69,16 @@ where
     I: Iterator<Item = AssetRaw>,
 {
     iter.map(|asset| -> Result<_> {
-        let agent_id = agent_ids.get_id_by_str(&asset.agent_id)?;
+        let agent_id = agent_ids.get_id(&asset.agent_id)?;
         let process = processes
             .get(asset.process_id.as_str())
             .with_context(|| format!("Invalid process ID: {}", &asset.process_id))?;
-        let region_id = region_ids.get_id_by_str(&asset.region_id)?;
+        let region_id = region_ids.get_id(&asset.region_id)?;
 
         Asset::new(
-            agent_id,
+            agent_id.clone(),
             Rc::clone(process),
-            region_id,
+            region_id.clone(),
             asset.capacity,
             asset.commission_year,
         )

src/input/commodity/cost.rs

@@ -78,7 +78,7 @@ where
     let mut commodity_regions: HashMap<CommodityID, HashSet<RegionID>> = HashMap::new();
 
     for cost in iter {
-        let commodity_id = commodity_ids.get_id_by_str(&cost.commodity_id)?;
+        let commodity_id = commodity_ids.get_id(&cost.commodity_id)?;
         let regions = parse_region_str(&cost.regions, region_ids)?;
         let years = parse_year_str(&cost.years, milestone_years)?;
         let ts_selection = time_slice_info.get_selection(&cost.time_slice)?;

src/input/commodity/demand.rs

@@ -110,14 +110,14 @@ where
     let mut map = AnnualDemandMap::new();
     for demand in iter {
         let commodity_id = svd_commodity_ids
-            .get_id_by_str(&demand.commodity_id)
+            .get_id(&demand.commodity_id)
             .with_context(|| {
                 format!(
                     "Can only provide demand data for SVD commodities. Found entry for '{}'",
                     demand.commodity_id
                 )
             })?;
-        let region_id = region_ids.get_id_by_str(&demand.region_id)?;
+        let region_id = region_ids.get_id(&demand.region_id)?;
 
         ensure!(
             milestone_years.binary_search(&demand.year).is_ok(),

src/input/commodity/demand_slicing.rs

@@ -64,14 +64,14 @@ where
 
     for slice in iter {
         let commodity_id = svd_commodity_ids
-            .get_id_by_str(&slice.commodity_id)
+            .get_id(&slice.commodity_id)
             .with_context(|| {
                 format!(
                     "Can only provide demand slice data for SVD commodities. Found entry for '{}'",
                     slice.commodity_id
                 )
             })?;
-        let region_id = region_ids.get_id_by_str(&slice.region_id)?;
+        let region_id = region_ids.get_id(&slice.region_id)?;
 
         // We need to know how many time slices are covered by the current demand slice entry and
         // how long they are relative to one another so that we can divide up the demand for this

src/input/process/availability.rs

@@ -109,9 +109,9 @@ where
         record.validate()?;
 
         // Get process
-        let id = process_ids.get_id_by_str(&record.process_id)?;
+        let id = process_ids.get_id(&record.process_id)?;
         let process = processes
-            .get(&id)
+            .get(id)
             .with_context(|| format!("Process {id} not found"))?;
 
         // Get regions
@@ -131,7 +131,9 @@ where
         let ts_selection = time_slice_info.get_selection(&record.time_slice)?;
 
         // Insert the energy limit into the map
-        let entry = map.entry(id).or_insert_with(ProcessEnergyLimitsMap::new);
+        let entry = map
+            .entry(id.clone())
+            .or_insert_with(ProcessEnergyLimitsMap::new);
         for (time_slice, ts_length) in time_slice_info.iter_selection(&ts_selection) {
             let bounds = record.to_bounds(ts_length);
 

src/input/process/flow.rs

@@ -82,9 +82,9 @@ where
         record.validate()?;
 
         // Get process
-        let id = process_ids.get_id_by_str(&record.process_id)?;
+        let id = process_ids.get_id(&record.process_id)?;
         let process = processes
-            .get(&id)
+            .get(id)
             .with_context(|| format!("Process {id} not found"))?;
 
         // Get regions