program: use sparse matrix for constituent map and update tests

Author: 0xbigz

programs/drift/src/state/lp_pool.rs

@@ -134,18 +134,24 @@ pub struct Constituent {
 //   pub struct PerpConstituent {
 //   }
 
+pub struct AmmConstituentDatum {
+    pub perp_market_index: u16,
+    pub constituent_index: u16,
+		/// PERCENTAGE_PRECISION. The weight this constituent has on the perp market
+    pub data: u64,
+    pub last_slot: u64,
+}
+
 pub struct WeightDatum {
+    pub constituent_index: u16,
+		/// PERCENTAGE_PRECISION. The weights of the target weight matrix
     pub data: u64,
     pub last_slot: u64,
 }
 
 pub struct AmmConstituentMapping {
-    // rows in the matrix, (perp markets)
-    pub num_rows: u16,
-    // columns in the matrix (VaultConstituents, spot markets)
-    pub num_cols: u16,
     // flattened matrix elements, PERCENTAGE_PRECISION. Keep at the end of the account to allow expansion with new constituents.
-    pub data: Vec<WeightDatum>,
+    pub data: Vec<AmmConstituentDatum>,
 }
 
 pub struct ConstituentTargetWeights {
@@ -175,15 +181,13 @@ impl ConstituentTargetWeights {
     pub fn update_target_weights(
         &mut self,
         mapping: &AmmConstituentMapping,
-        amm_inventory: &[u64], // length = mapping.num_rows
+        amm_inventory: &[(u16, u64)], // length = mapping.num_rows
         constituents: &[Constituent],
         prices: &[u64], // same order as constituents
         aum: u64,
         slot: u64,
     ) -> DriftResult<()> {
         // assert_ne!(aum, 0);
-        assert_eq!(constituents.len(), mapping.num_cols as usize);
-        assert_eq!(amm_inventory.len(), mapping.num_rows as usize);
         assert_eq!(prices.len(), constituents.len());
 
         self.data.clear();
@@ -194,11 +198,13 @@ impl ConstituentTargetWeights {
         for (constituent_index, constituent) in constituents.iter().enumerate() {
             let mut target_amount = 0u128;
 
-            for (row_index, &inventory) in amm_inventory.iter().enumerate() {
-                let idx = row_index * mapping.num_cols as usize + constituent_index;
+            for (perp_market_index, inventory) in amm_inventory.iter() {
+                let idx = mapping.data
+                    .iter()
+                    .position(|d| &d.perp_market_index == perp_market_index)
+                    .expect("missing mapping for this market index");
                 let weight = mapping.data[idx].data as u128; // PERCENTAGE_PRECISION
-
-                target_amount += inventory as u128 * weight / PERCENTAGE_PRECISION_U64 as u128;
+                target_amount += (*inventory) as u128 * weight / PERCENTAGE_PRECISION_U64 as u128;
             }
 
             let price = prices[constituent_index] as u128;
@@ -210,6 +216,7 @@ impl ConstituentTargetWeights {
             let weight_datum = (target_weight as u64).min(PERCENTAGE_PRECISION_U64);
 
             self.data.push(WeightDatum {
+                constituent_index: constituent_index as u16,
                 data: weight_datum,
                 last_slot: slot,
             });

programs/drift/src/state/lp_pool/tests.rs

@@ -8,8 +8,11 @@ mod tests {
 
     const PERCENTAGE_PRECISION_U64: u64 = 1_000_000;
 
-    fn weight_datum(data: u64, last_slot: u64) -> WeightDatum {
-        WeightDatum { data, last_slot }
+    fn weight_datum(constituent_index: u16, data: u64, last_slot: u64) -> WeightDatum {
+        WeightDatum { constituent_index, data, last_slot }
+    }
+    fn amm_const_datum(perp_market_index: u16, constituent_index: u16, data: u64, last_slot: u64) -> AmmConstituentDatum {
+        AmmConstituentDatum { perp_market_index, constituent_index, data, last_slot }
     }
 
     fn dummy_constituent(index: u16) -> Constituent {
@@ -37,12 +40,10 @@ mod tests {
     #[test]
     fn test_single_zero_weight() {
         let mapping = AmmConstituentMapping {
-            num_rows: 1,
-            num_cols: 1,
-            data: vec![weight_datum(0, 0)],
+            data: vec![amm_const_datum(0, 1, 0, 0)],
         };
 
-        let amm_inventory = vec![1_000_000];
+        let amm_inventory: Vec<(u16, u64)> = vec![(0, 1_000_000)];
         let prices = vec![1_000_000];
         let constituents = vec![dummy_constituent(0)];
         let aum = 1_000_000;
@@ -68,12 +69,10 @@ mod tests {
     #[test]
     fn test_single_full_weight() {
         let mapping = AmmConstituentMapping {
-            num_rows: 1,
-            num_cols: 1,
-            data: vec![weight_datum(PERCENTAGE_PRECISION_U64, 0)],
+            data: vec![amm_const_datum(0, 1, PERCENTAGE_PRECISION_U64, 0)],
         };
 
-        let amm_inventory = vec![1_000_000];
+        let amm_inventory = vec![(0, 1_000_000)];
         let prices = vec![1_000_000];
         let constituents = vec![dummy_constituent(0)];
         let aum = 1_000_000;
@@ -99,15 +98,13 @@ mod tests {
     #[test]
     fn test_multiple_constituents_partial_weights() {
         let mapping = AmmConstituentMapping {
-            num_rows: 1,
-            num_cols: 2,
             data: vec![
-                weight_datum(PERCENTAGE_PRECISION_U64 / 2, 0),
-                weight_datum(PERCENTAGE_PRECISION_U64 / 2, 0),
+                amm_const_datum(0, 1, PERCENTAGE_PRECISION_U64 / 2, 0),
+                amm_const_datum(0, 2, PERCENTAGE_PRECISION_U64 / 2, 0),
             ],
         };
 
-        let amm_inventory = vec![1_000_000];
+        let amm_inventory = vec![(0, 1_000_000)];
         let prices = vec![1_000_000, 1_000_000];
         let constituents = vec![dummy_constituent(0), dummy_constituent(1)];
         let aum = 1_000_000;
@@ -136,12 +133,10 @@ mod tests {
     #[test]
     fn test_zero_aum_safe() {
         let mapping = AmmConstituentMapping {
-            num_rows: 1,
-            num_cols: 1,
-            data: vec![weight_datum(PERCENTAGE_PRECISION_U64, 0)],
+            data: vec![amm_const_datum(0, 1, PERCENTAGE_PRECISION_U64, 0)],
         };
 
-        let amm_inventory = vec![1_000_000];
+        let amm_inventory = vec![(0, 1_000_000)];
         let prices = vec![1_000_000];
         let constituents = vec![dummy_constituent(0)];
         let aum = 0;
@@ -167,12 +162,10 @@ mod tests {
     #[test]
     fn test_overflow_protection() {
         let mapping = AmmConstituentMapping {
-            num_rows: 1,
-            num_cols: 1,
-            data: vec![weight_datum(u64::MAX, 0)],
+            data: vec![amm_const_datum(0, 1, u64::MAX, 0)],
         };
 
-        let amm_inventory = vec![u64::MAX];
+        let amm_inventory = vec![(0, u64::MAX)];
         let prices = vec![u64::MAX];
         let constituents = vec![dummy_constituent(0)];
         let aum = 1;