lockup: Don't push back vesting window when calculating unlock

.travis.yml

@@ -50,4 +50,5 @@ jobs:
     - <<: *defaults
       name: Runs the tests
       script:
+        - cargo test --lib
         - anchor test

programs/lockup/src/calculator.rs

@@ -38,47 +38,143 @@ fn total_vested(vesting: &Vesting, current_ts: i64) -> u64 {
     } else if current_ts >= vesting.end_ts {
         vesting.start_balance
     } else {
-        linear_unlock(vesting, current_ts).unwrap()
+        linear_unlock(vesting, current_ts)
     }
 }
 
-fn linear_unlock(vesting: &Vesting, current_ts: i64) -> Option<u64> {
+// Assumes `current_ts` < `vesting.end_ts`.
+fn linear_unlock(vesting: &Vesting, current_ts: i64) -> u64 {
     // Signed division not supported.
-    let current_ts = current_ts as u64;
-    let start_ts = vesting.start_ts as u64;
-    let end_ts = vesting.end_ts as u64;
+    let current_ts = current_ts as f64;
+    let start_ts = vesting.start_ts as f64;
+    let end_ts = vesting.end_ts as f64;
 
-    // If we can't perfectly partition the vesting window,
-    // push the start of the window back so that we can.
-    //
-    // This has the effect of making the first vesting period shorter
-    // than the rest.
-    let shifted_start_ts =
-        start_ts.checked_sub(end_ts.checked_sub(start_ts)? % vesting.period_count)?;
+    // The length of a single vesting period.
+    // Invariant: period_count <= (end_ts - start_ts).
+    let period_secs: f64 = (end_ts - start_ts) / (vesting.period_count as f64);
 
-    // Similarly, if we can't perfectly divide up the vesting rewards
-    // then make the first period act as a cliff, earning slightly more than
-    // subsequent periods.
-    let reward_overflow = vesting.start_balance % vesting.period_count;
+    // The period the current_ts is in (floor divides).
+    // Invariant: current_ts >= start_ts.
+    let current_period: u64 = ((current_ts - start_ts) / period_secs) as u64;
 
-    // Reward per period ignoring the overflow.
-    let reward_per_period =
-        (vesting.start_balance.checked_sub(reward_overflow)?).checked_div(vesting.period_count)?;
+    // Reward per period.
+    let reward_per_period: f64 = (vesting.start_balance as f64) / (vesting.period_count as f64);
 
-    // Number of vesting periods that have passed.
-    let current_period = {
-        let period_secs =
-            (end_ts.checked_sub(shifted_start_ts)?).checked_div(vesting.period_count)?;
-        let current_period_count =
-            (current_ts.checked_sub(shifted_start_ts)?).checked_div(period_secs)?;
-        std::cmp::min(current_period_count, vesting.period_count)
-    };
+    // Rounds the total reward down to the nearest integer, since we can't
+    // pay out fractional rewards.
+    ((current_period as f64) * reward_per_period) as u64
+}
 
-    if current_period == 0 {
-        return Some(0);
+#[cfg(test)]
+mod tests {
+    use super::*;
+    use anchor_lang::solana_program::pubkey::Pubkey;
+
+    // Window = 10 seconds.
+    // Period count = 2.
+    // =>
+    // Every 5 seconds 2.5 is vested.
+    #[test]
+    fn vesting_window_evenly_divisible_by_period_count() {
+        let v = create_vesting(5, 10, 20, 2);
+        let cases = vec![
+            [0, 0], // Before vesting begins.
+            [9, 0],
+            [10, 0], // Vesting begins.
+            [11, 0],
+            [12, 0],
+            [13, 0],
+            [14, 0],
+            [15, 2], // 2.5 is vested (floor).
+            [16, 2],
+            [17, 2],
+            [18, 2],
+            [19, 2],
+            [20, 5], // All vested.
+            [21, 5],
+        ];
+        run_test(v, cases);
     }
 
-    current_period
-        .checked_mul(reward_per_period)?
-        .checked_add(reward_overflow)
+    // Window = 11 seconds.
+    // Period count = 2.
+    // =>
+    // Every 5.5 seconds 2.5 is vested.
+    #[test]
+    fn vesting_window_not_evenly_divisble_by_period_count() {
+        let v = create_vesting(5, 10, 21, 2);
+        let cases = vec![
+            [10, 0], // Vesting begins.
+            [11, 0],
+            [12, 0],
+            [13, 0],
+            [14, 0],
+            [15, 0],
+            [16, 2], // 2.5 vested.
+            [17, 2],
+            [18, 2],
+            [19, 2],
+            [20, 2],
+            [21, 5], // All vested.
+            [22, 5],
+        ];
+        run_test(v, cases);
+    }
+
+    // Winow = 11 seconds.
+    // Period_count = 6.
+    // =>
+    // Every 1.83 seconds about 16.67 is vested.
+    #[test]
+    fn cumulative_remainder() {
+        let v = create_vesting(100, 30, 41, 6);
+        let cases = vec![
+            [30, 0], // Vesting begins.
+            [31, 0],
+            [32, 16], // 16.67 @ 1.83 seconds.
+            [33, 16],
+            [34, 33], // 33.34 @ 3.66 seconds.
+            [35, 33],
+            [36, 50], // 50.01 @ 5.49 seconds.
+            [37, 50],
+            [38, 66], // 66.68 @ 7.32 seconds.
+            [39, 66],
+            [40, 83],  // 83.35 @ 9.15 seconds.
+            [41, 100], // 100 @ 11 seconds.
+        ];
+        run_test(v, cases);
+    }
+
+    // Each case is an array consisting of
+    // [start_balance, start_ts, end_ts, period_count, current_ts, total_vested].
+    fn run_test(v: Vesting, cases: Vec<[u64; 2]>) {
+        for c in cases.iter() {
+            println!("Case: {:?}", c);
+            let r = total_vested(&v, c[0] as i64);
+            assert_eq!(r, c[1])
+        }
+    }
+
+    fn create_vesting(
+        start_balance: u64,
+        start_ts: i64,
+        end_ts: i64,
+        period_count: u64,
+    ) -> Vesting {
+        Vesting {
+            beneficiary: Pubkey::new_unique(),
+            mint: Pubkey::new_unique(),
+            vault: Pubkey::new_unique(),
+            grantor: Pubkey::new_unique(),
+            outstanding: 0,
+            start_balance,
+            created_ts: 0,
+            start_ts,
+            end_ts,
+            period_count,
+            whitelist_owned: 0,
+            nonce: 0,
+            realizor: None,
+        }
+    }
 }