fix a couple nits with the interfaces and code

README.md

@@ -26,8 +26,8 @@ Contracts in this repository are designed so that they may be used together _or_
 
 `Staker` enables users to delegate the balance of a token towards an account, and tracks the historical delegation at each block. In addition, it allows the computation of the time-weighted average delegated tokens of any account over any historical period.
 
-- Users call `Token#approve(staker, stake_amount)`, then `Staker#stake(delegate)` to stake and delegate their tokens to other addresses
-- Users call `Staker#withdraw(delegate, recipient, amount)` to remove part or all of their delegation
+- Users call `Token#approve(staker, stake_amount)`, then `Staker#stake_amount(delegate, stake_amount)` to stake and delegate their tokens to other addresses
+- Users call `Staker#withdraw_amount(delegate, recipient, amount)` to remove part or all of their delegation
 - The average delegation weight is computable over *any* historical period
 - The contract has no owner, and cannot be updated nor configured
 

src/governor.cairo

@@ -43,10 +43,15 @@ pub trait IGovernor<TContractState> {
     // Vote on the given proposal.
     fn vote(ref self: TContractState, id: felt252, yea: bool);
 
-    // Cancel the given proposal. The proposer may cancel the proposal at any time during before or during the voting period.
-    // Cancellation can happen by any address if the average voting weight is below the proposal_creation_threshold.
+    // Cancel the proposal with the given ID. Same as #cancel_at_timestamp, but uses the current timestamp for computing the voting weight.
     fn cancel(ref self: TContractState, id: felt252);
 
+    // Cancel the proposal with the given ID. The proposal may be canceled at any time before it is executed.
+    // There are two ways the proposal cancellation can be authorized:
+    // - The proposer can cancel the proposal
+    // - Anyone can cancel if the average voting weight of the proposer was below the proposal_creation_threshold during the voting period (at the given breach_timestamp)
+    fn cancel_at_timestamp(ref self: TContractState, id: felt252, breach_timestamp: u64);
+
     // Execute the given proposal.
     fn execute(ref self: TContractState, call: Call) -> Span<felt252>;
 
@@ -100,6 +105,7 @@ pub mod Governor {
     #[derive(starknet::Event, Drop)]
     pub struct Canceled {
         pub id: felt252,
+        pub breach_timestamp: u64,
     }
 
     #[derive(starknet::Event, Drop)]
@@ -150,7 +156,8 @@ pub mod Governor {
             if latest_proposal_id.is_non_zero() {
                 let latest_proposal_state = self.get_proposal(latest_proposal_id).execution_state;
 
-                if (latest_proposal_state.canceled.is_zero()) {
+                // if the proposal is not canceled, check that the voting for that proposal has ended
+                if latest_proposal_state.canceled.is_zero() {
                     assert(
                         latest_proposal_state.created
                             + config.voting_start_delay
@@ -236,54 +243,60 @@ pub mod Governor {
             self.proposals.write(id, proposal);
             self.has_voted.write((voter, id), true);
 
-            self.emit(Voted { id, voter, weight, yea, });
+            self.emit(Voted { id, voter, weight, yea });
         }
 
 
         fn cancel(ref self: ContractState, id: felt252) {
+            self.cancel_at_timestamp(id, get_block_timestamp())
+        }
+
+        fn cancel_at_timestamp(ref self: ContractState, id: felt252, breach_timestamp: u64) {
             let config = self.config.read();
-            let voting_token = self.staker.read();
+            let staker = self.staker.read();
             let mut proposal = self.proposals.read(id);
 
             assert(proposal.proposer.is_non_zero(), 'DOES_NOT_EXIST');
 
-            if (proposal.proposer != get_caller_address()) {
-                // if at any point the average voting weight is below the proposal_creation_threshold for the proposer, it can be canceled
+            assert(proposal.execution_state.canceled.is_zero(), 'ALREADY_CANCELED');
+            assert(proposal.execution_state.executed.is_zero(), 'ALREADY_EXECUTED');
+            assert(breach_timestamp >= proposal.execution_state.created, 'PROPOSAL_NOT_CREATED');
+
+            assert(
+                breach_timestamp < (proposal.execution_state.created
+                    + config.voting_start_delay
+                    + config.voting_period),
+                'VOTING_ENDED'
+            );
+
+            // iff the proposer is not calling this we need to check the voting weight
+            if proposal.proposer != get_caller_address() {
+                // if at the given timestamp (during the voting period),
+                // the average voting weight is below the proposal_creation_threshold for the proposer, it can be canceled
                 assert(
-                    voting_token
-                        .get_average_delegated_over_last(
+                    staker
+                        .get_average_delegated(
                             delegate: proposal.proposer,
-                            period: config.voting_weight_smoothing_duration
+                            start: breach_timestamp - config.voting_weight_smoothing_duration,
+                            end: breach_timestamp
                         ) < config
                         .proposal_creation_threshold,
                     'THRESHOLD_NOT_BREACHED'
                 );
             }
 
-            let timestamp_current = get_block_timestamp();
-
-            assert(
-                timestamp_current < (proposal.execution_state.created
-                    + config.voting_start_delay
-                    + config.voting_period),
-                'VOTING_ENDED'
-            );
-
-            // we know it's not executed since we check voting has not ended
             proposal
                 .execution_state =
                     ExecutionState {
                         created: proposal.execution_state.created,
+                        // we asserted that it is not already executed
                         executed: 0,
-                        canceled: timestamp_current
+                        canceled: get_block_timestamp()
                     };
 
             self.proposals.write(id, proposal);
 
-            // allows the proposer to create a new proposal
-            self.latest_proposal_by_proposer.write(proposal.proposer, Zero::zero());
-
-            self.emit(Canceled { id });
+            self.emit(Canceled { id, breach_timestamp });
         }
 
         fn execute(ref self: ContractState, call: Call) -> Span<felt252> {

src/governor_test.cairo

@@ -424,6 +424,14 @@ fn test_vote_after_voting_period_should_fail() {
     governor.vote(id, true); // vote should fail
 }
 
+#[test]
+#[should_panic(expected: ('DOES_NOT_EXIST', 'ENTRYPOINT_FAILED'))]
+fn test_cancel_fails_if_proposal_not_exists() {
+    let (_staker, _token, governor, _config) = setup();
+    let id = 1234;
+    governor.cancel(id);
+}
+
 #[test]
 fn test_cancel_by_proposer() {
     let (staker, token, governor, config) = setup();
@@ -453,12 +461,28 @@ fn test_cancel_by_proposer() {
     );
 }
 
+#[test]
+#[should_panic(expected: ('ALREADY_CANCELED', 'ENTRYPOINT_FAILED'))]
+fn test_double_cancel_by_proposer() {
+    let (staker, token, governor, _config) = setup();
+
+    let proposer = proposer();
+
+    let id = create_proposal(governor, token, staker);
+
+    advance_time(30);
+
+    set_contract_address(proposer);
+    governor.cancel(id);
+    governor.cancel(id);
+}
+
 #[test]
 fn test_cancel_by_non_proposer() {
     let (staker, token, governor, config) = setup();
 
     let id = create_proposal(governor, token, staker);
-    staker.withdraw(proposer(), recipient: Zero::zero(), amount: 25);
+    staker.withdraw_amount(proposer(), recipient: Zero::zero(), amount: 25);
 
     // Fast forward one smoothing duration
     advance_time(config.voting_weight_smoothing_duration);
@@ -630,7 +654,7 @@ fn test_verify_votes_are_counted_over_voting_weight_smoothing_duration_from_star
     advance_time(config.voting_start_delay - (config.voting_weight_smoothing_duration / 3));
     // undelegate 1/3rd of a duration before voting starts, so only a third of voting power is counted for voter1
     set_contract_address(voter1);
-    staker.withdraw(voter1, recipient: Zero::zero(), amount: 49);
+    staker.withdraw_amount(voter1, recipient: Zero::zero(), amount: 49);
 
     advance_time((config.voting_weight_smoothing_duration / 3));
 

src/staker.cairo

@@ -13,8 +13,14 @@ pub trait IStaker<TContractState> {
     // Transfer the approved amount of token from the caller into this contract and delegates it to the given address
     fn stake(ref self: TContractState, delegate: ContractAddress);
 
-    // Withdraws the delegated tokens from the contract to the given recipient address
-    fn withdraw(
+    // Transfer the specified amount of token from the caller into this contract and delegates the voting weight to the specified delegate
+    fn stake_amount(ref self: TContractState, delegate: ContractAddress, amount: u128);
+
+    // Unstakes and withdraws all of the tokens delegated by the sender to the delegate from the contract to the given recipient address
+    fn withdraw(ref self: TContractState, delegate: ContractAddress, recipient: ContractAddress);
+
+    // Unstakes and withdraws the specified amount of tokens delegated by the sender to the delegate from the contract to the given recipient address
+    fn withdraw_amount(
         ref self: TContractState,
         delegate: ContractAddress,
         recipient: ContractAddress,
@@ -214,27 +220,45 @@ pub mod Staker {
         }
 
         fn stake(ref self: ContractState, delegate: ContractAddress) {
+            self
+                .stake_amount(
+                    delegate,
+                    self
+                        .token
+                        .read()
+                        .allowance(get_caller_address(), get_contract_address())
+                        .try_into()
+                        .expect('ALLOWANCE_OVERFLOW')
+                );
+        }
+
+        fn stake_amount(ref self: ContractState, delegate: ContractAddress, amount: u128) {
             let from = get_caller_address();
-            let this_address = get_contract_address();
             let token = self.token.read();
-            let amount = token.allowance(from, this_address);
 
-            let amount_small: u128 = amount.try_into().expect('ALLOWANCE_OVERFLOW');
             assert(
-                token.transferFrom(from, get_contract_address(), amount), 'TRANSFER_FROM_FAILED'
+                token.transferFrom(from, get_contract_address(), amount.into()),
+                'TRANSFER_FROM_FAILED'
             );
 
             let key = (from, delegate);
-            self.staked.write((from, delegate), amount_small + self.staked.read(key));
+            self.staked.write((from, delegate), amount + self.staked.read(key));
             self
                 .amount_delegated
-                .write(
-                    delegate, self.insert_snapshot(delegate, get_block_timestamp()) + amount_small
-                );
-            self.emit(Staked { from, delegate, amount: amount_small });
+                .write(delegate, self.insert_snapshot(delegate, get_block_timestamp()) + amount);
+            self.emit(Staked { from, delegate, amount });
         }
 
         fn withdraw(
+            ref self: ContractState, delegate: ContractAddress, recipient: ContractAddress
+        ) {
+            self
+                .withdraw_amount(
+                    delegate, recipient, self.staked.read((get_caller_address(), delegate))
+                )
+        }
+
+        fn withdraw_amount(
             ref self: ContractState,
             delegate: ContractAddress,
             recipient: ContractAddress,

src/staker_test.cairo

@@ -305,7 +305,7 @@ fn test_get_average_delegated() {
 
     // rewind to undelegate at 8
     set_block_timestamp(8);
-    staker.withdraw(delegatee, recipient: contract_address_const::<0>(), amount: 12345);
+    staker.withdraw_amount(delegatee, recipient: contract_address_const::<0>(), amount: 12345);
 
     set_block_timestamp(12);
     assert(staker.get_average_delegated(delegatee, 4, 10) == 8230, 'average (4 sec * 12345)/6');
@@ -320,7 +320,7 @@ fn test_transfer_delegates_moved() {
     set_block_timestamp(2);
     token.approve(staker.contract_address, 12345);
     staker.stake(delegatee);
-    staker.withdraw(delegatee, contract_address_const::<3456>(), 500);
+    staker.withdraw_amount(delegatee, contract_address_const::<3456>(), 500);
     set_block_timestamp(5);
 
     assert_eq!(staker.get_delegated(delegatee), (12345 - 500));
@@ -338,7 +338,7 @@ fn test_delegate_undelegate() {
     staker.stake(delegatee);
 
     set_block_timestamp(5);
-    staker.withdraw(delegatee, Zero::zero(), 12345);
+    staker.withdraw_amount(delegatee, Zero::zero(), 12345);
     set_block_timestamp(8);
 
     assert(staker.get_delegated(delegatee) == 0, 'delegated');