feat: improve FinalityProofProvider api

client/finality-grandpa/src/finality_proof.rs

@@ -114,7 +114,8 @@ where
 			return Ok(None)
 		};
 
-		prove_finality(&*self.backend, authority_set_changes, block)
+		prove_finality(&*self.backend, authority_set_changes, block, true)
+			.map(|opt| opt.map(|proof| proof.encode()))
 	}
 }
 
@@ -146,11 +147,14 @@ pub enum FinalityProofError {
 	Client(#[from] sp_blockchain::Error),
 }
 
-fn prove_finality<Block, B>(
+/// Prove finality for the given block number by returning a Justification for the last block of
+/// the authority set.
+pub fn prove_finality<Block, B>(
 	backend: &B,
 	authority_set_changes: AuthoritySetChanges<NumberFor<Block>>,
 	block: NumberFor<Block>,
-) -> Result<Option<Vec<u8>>, FinalityProofError>
+	collect_unknown_headers: bool,
+) -> Result<Option<FinalityProof<Block::Header>>, FinalityProofError>
 where
 	Block: BlockT,
 	B: Backend<Block>,
@@ -214,8 +218,8 @@ where
 		},
 	};
 
-	// Collect all headers from the requested block until the last block of the set
-	let unknown_headers = {
+	let unknown_headers = if collect_unknown_headers {
+		// Collect all headers from the requested block until the last block of the set
 		let mut headers = Vec::new();
 		let mut current = block + One::one();
 		loop {
@@ -227,16 +231,15 @@ where
 			current += One::one();
 		}
 		headers
+	} else {
+		Default::default()
 	};
 
-	Ok(Some(
-		FinalityProof {
-			block: backend.blockchain().expect_block_hash_from_id(&BlockId::Number(just_block))?,
-			justification,
-			unknown_headers,
-		}
-		.encode(),
-	))
+	Ok(Some(FinalityProof {
+		block: backend.blockchain().expect_block_hash_from_id(&BlockId::Number(just_block))?,
+		justification,
+		unknown_headers,
+	}))
 }
 
 #[cfg(test)]
@@ -334,7 +337,7 @@ mod tests {
 		let authority_set_changes = AuthoritySetChanges::empty();
 
 		// The last finalized block is 4, so we cannot provide further justifications.
-		let proof_of_5 = prove_finality(&*backend, authority_set_changes, 5);
+		let proof_of_5 = prove_finality(&*backend, authority_set_changes, 5, true);
 		assert!(matches!(proof_of_5, Err(FinalityProofError::BlockNotYetFinalized)));
 	}
 
@@ -347,7 +350,7 @@ mod tests {
 
 		// Block 4 is finalized without justification
 		// => we can't prove finality of 3
-		let proof_of_3 = prove_finality(&*backend, authority_set_changes, 3).unwrap();
+		let proof_of_3 = prove_finality(&*backend, authority_set_changes, 3, true).unwrap();
 		assert_eq!(proof_of_3, None);
 	}
 
@@ -478,7 +481,7 @@ mod tests {
 		let mut authority_set_changes = AuthoritySetChanges::empty();
 		authority_set_changes.append(1, 8);
 
-		let proof_of_6 = prove_finality(&*backend, authority_set_changes, 6);
+		let proof_of_6 = prove_finality(&*backend, authority_set_changes, 6, true);
 		assert!(matches!(proof_of_6, Err(FinalityProofError::BlockNotInAuthoritySetChanges)));
 	}
 
@@ -502,10 +505,11 @@ mod tests {
 		authority_set_changes.append(0, 5);
 		authority_set_changes.append(1, 8);
 
-		let proof_of_6: FinalityProof = Decode::decode(
-			&mut &prove_finality(&*backend, authority_set_changes.clone(), 6).unwrap().unwrap()[..],
-		)
-		.unwrap();
+		let proof_of_6: FinalityProof =
+			prove_finality(&*backend, authority_set_changes.clone(), 6, true)
+				.unwrap()
+				.unwrap();
+
 		assert_eq!(
 			proof_of_6,
 			FinalityProof {
@@ -514,6 +518,20 @@ mod tests {
 				unknown_headers: vec![block7.header().clone(), block8.header().clone()],
 			},
 		);
+
+		let proof_of_6_without_unknown: FinalityProof =
+			prove_finality(&*backend, authority_set_changes.clone(), 6, false)
+				.unwrap()
+				.unwrap();
+
+		assert_eq!(
+			proof_of_6_without_unknown,
+			FinalityProof {
+				block: block8.hash(),
+				justification: grandpa_just8.encode(),
+				unknown_headers: vec![],
+			},
+		);
 	}
 
 	#[test]
@@ -525,7 +543,7 @@ mod tests {
 		let mut authority_set_changes = AuthoritySetChanges::empty();
 		authority_set_changes.append(0, 5);
 
-		assert!(matches!(prove_finality(&*backend, authority_set_changes, 6), Ok(None)));
+		assert!(matches!(prove_finality(&*backend, authority_set_changes, 6, true), Ok(None)));
 	}
 
 	#[test]
@@ -544,10 +562,9 @@ mod tests {
 		let mut authority_set_changes = AuthoritySetChanges::empty();
 		authority_set_changes.append(0, 5);
 
-		let proof_of_6: FinalityProof = Decode::decode(
-			&mut &prove_finality(&*backend, authority_set_changes, 6).unwrap().unwrap()[..],
-		)
-		.unwrap();
+		let proof_of_6: FinalityProof =
+			prove_finality(&*backend, authority_set_changes, 6, true).unwrap().unwrap();
+
 		assert_eq!(
 			proof_of_6,
 			FinalityProof {

client/finality-grandpa/src/lib.rs

@@ -130,7 +130,7 @@ pub use authorities::{AuthoritySet, AuthoritySetChanges, SharedAuthoritySet};
 pub use aux_schema::best_justification;
 pub use communication::grandpa_protocol_name::standard_name as protocol_standard_name;
 pub use finality_grandpa::voter::report;
-pub use finality_proof::{FinalityProof, FinalityProofError, FinalityProofProvider};
+pub use finality_proof::{FinalityProof, FinalityProofError, FinalityProofProvider, prove_finality};
 pub use import::{find_forced_change, find_scheduled_change, GrandpaBlockImport};
 pub use justification::GrandpaJustification;
 pub use notification::{GrandpaJustificationSender, GrandpaJustificationStream};