1  // Copyright (c) 2025 ADnet Contributors
  2  // This file is part of the AlphaVM library.
  3  
  4  // Licensed under the Apache License, Version 2.0 (the "License");
  5  // you may not use this file except in compliance with the License.
  6  // You may obtain a copy of the License at:
  7  
  8  // http://www.apache.org/licenses/LICENSE-2.0
  9  
 10  // Unless required by applicable law or agreed to in writing, software
 11  // distributed under the License is distributed on an "AS IS" BASIS,
 12  // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 13  // See the License for the specific language governing permissions and
 14  // limitations under the License.
 15  
 16  use crate::{
 17      prelude::{FromBytes, Identifier, IoResult, Network, Read, ToBytes},
 18      synthesizer::{Program, snark::ProvingKey},
 19  };
 20  
 21  use anyhow::{Result, anyhow, bail, ensure};
 22  use std::{
 23      fs::{self, File},
 24      io::Write,
 25      path::Path,
 26  };
 27  
 28  static PROVER_FILE_EXTENSION: &str = "prover";
 29  
 30  pub struct ProverFile<N: Network> {
 31      /// The function name.
 32      function_name: Identifier<N>,
 33      /// The proving key.
 34      proving_key: ProvingKey<N>,
 35  }
 36  
 37  impl<N: Network> ProverFile<N> {
 38      /// Creates a new proving key file, given the directory path, function name, and proving key.
 39      pub fn create(directory: &Path, function_name: &Identifier<N>, proving_key: ProvingKey<N>) -> Result<Self> {
 40          // Ensure the directory path exists.
 41          ensure!(directory.exists(), "The build directory does not exist: '{}'", directory.display());
 42          // Ensure the function name is valid.
 43          ensure!(!Program::is_reserved_keyword(function_name), "Function name is invalid (reserved): {function_name}");
 44  
 45          // Create the candidate prover file.
 46          let prover_file = Self { function_name: *function_name, proving_key };
 47  
 48          // Create the file name.
 49          let file_name = format!("{function_name}.{PROVER_FILE_EXTENSION}");
 50          // Construct the file path.
 51          let path = directory.join(file_name);
 52          // Write the file (overwriting if it already exists).
 53          File::create(&path)?.write_all(&prover_file.to_bytes_le()?)?;
 54  
 55          // Attempt to load the prover file.
 56          Self::from_filepath(&path)
 57      }
 58  
 59      /// Opens the prover file, given the directory path and function name.
 60      pub fn open(directory: &Path, function_name: &Identifier<N>) -> Result<Self> {
 61          // Ensure the directory path exists.
 62          ensure!(directory.exists(), "The build directory does not exist: '{}'", directory.display());
 63  
 64          // Create the file name.
 65          let file_name = format!("{function_name}.{PROVER_FILE_EXTENSION}");
 66          // Construct the file path.
 67          let path = directory.join(file_name);
 68          // Ensure the file path exists.
 69          ensure!(path.exists(), "The prover file is missing: '{}'", path.display());
 70  
 71          // Load the prover file.
 72          let prover = Self::from_filepath(&path)?;
 73  
 74          // Ensure the function name matches.
 75          if prover.function_name() != function_name {
 76              bail!(
 77                  "The prover file for '{}' contains an incorrect function name of '{}'",
 78                  function_name,
 79                  prover.function_name()
 80              );
 81          }
 82  
 83          Ok(prover)
 84      }
 85  
 86      /// Returns `true` if the prover file for the given function name exists at the given directory.
 87      pub fn exists_at(directory: &Path, function_name: &Identifier<N>) -> bool {
 88          // Create the file name.
 89          let file_name = format!("{function_name}.{PROVER_FILE_EXTENSION}");
 90          // Construct the file path.
 91          let path = directory.join(file_name);
 92          // Ensure the path is well-formed.
 93          Self::check_path(&path).is_ok() && path.exists()
 94      }
 95  
 96      /// Returns the function name.
 97      pub const fn function_name(&self) -> &Identifier<N> {
 98          &self.function_name
 99      }
100  
101      /// Returns the proving key.
102      pub const fn proving_key(&self) -> &ProvingKey<N> {
103          &self.proving_key
104      }
105  
106      /// Removes the file at the given path, if it exists.
107      pub fn remove(&self, path: &Path) -> Result<()> {
108          // If the path does not exist, do nothing.
109          if !path.exists() {
110              Ok(())
111          } else {
112              // Ensure the path is well-formed.
113              Self::check_path(path)?;
114              // Remove the file.
115              Ok(fs::remove_file(path)?)
116          }
117      }
118  }
119  
120  impl<N: Network> ProverFile<N> {
121      /// Checks that the given path has the correct file extension.
122      fn check_path(path: &Path) -> Result<()> {
123          // Ensure the given path is a file.
124          ensure!(path.is_file(), "The path is not a file.");
125  
126          // Ensure the given path has the correct file extension.
127          let extension = path.extension().ok_or_else(|| anyhow!("File extension not found."))?;
128          ensure!(extension == PROVER_FILE_EXTENSION, "File extension is incorrect.");
129  
130          // Ensure the given path exists.
131          ensure!(path.exists(), "File does not exist: {}", path.display());
132  
133          Ok(())
134      }
135  
136      /// Reads the prover from the given file path, if it exists.
137      fn from_filepath(file: &Path) -> Result<Self> {
138          // Ensure the path is well-formed.
139          Self::check_path(file)?;
140          // Parse the prover file bytes.
141          let prover = Self::from_bytes_le(&fs::read(file)?)?;
142  
143          // Retrieve the file stem.
144          let file_stem = file
145              .file_stem()
146              .ok_or_else(|| anyhow!("File name not found."))?
147              .to_str()
148              .ok_or_else(|| anyhow!("File name not found."))?
149              .to_string();
150          // Ensure the function name matches the file stem.
151          ensure!(prover.function_name.to_string() == file_stem, "Function name does not match file stem.");
152  
153          // Return the prover file.
154          Ok(prover)
155      }
156  
157      /// Writes the prover to the file.
158      pub fn write_to(&self, path: &Path) -> Result<()> {
159          // Ensure the path is well-formed.
160          Self::check_path(path)?;
161  
162          // Retrieve the file stem.
163          let file_stem = path
164              .file_name()
165              .ok_or_else(|| anyhow!("File name not found."))?
166              .to_str()
167              .ok_or_else(|| anyhow!("File name not found."))?
168              .to_string();
169          // Ensure the function name matches the file stem.
170          ensure!(self.function_name.to_string() == file_stem, "Function name does not match file stem.");
171  
172          // Write to the file (overwriting if it already exists).
173          Ok(File::create(path)?.write_all(&self.to_bytes_le()?)?)
174      }
175  }
176  
177  impl<N: Network> FromBytes for ProverFile<N> {
178      /// Reads the prover file from a buffer.
179      fn read_le<R: Read>(mut reader: R) -> IoResult<Self> {
180          let function_name = Identifier::read_le(&mut reader)?;
181          let proving_key = FromBytes::read_le(&mut reader)?;
182          Ok(Self { function_name, proving_key })
183      }
184  }
185  
186  impl<N: Network> ToBytes for ProverFile<N> {
187      /// Writes the prover file to a buffer.
188      fn write_le<W: Write>(&self, mut writer: W) -> IoResult<()> {
189          self.function_name.write_le(&mut writer)?;
190          self.proving_key.write_le(&mut writer)
191      }
192  }
193  
194  #[cfg(test)]
195  mod tests {
196      use super::*;
197      use crate::{
198          prelude::{FromStr, Parser, TestRng},
199          synthesizer::Process,
200      };
201  
202      type CurrentNetwork = alphavm_console::network::MainnetV0;
203      type CurrentAleo = alphavm_circuit::AleoV0;
204  
205      fn temp_dir() -> std::path::PathBuf {
206          tempfile::tempdir().expect("Failed to open temporary directory").keep()
207      }
208  
209      #[test]
210      #[ignore] // TODO: Requires credits.alpha → credits.alpha migration
211      fn test_create_and_open() {
212          // Initialize a temporary directory.
213          let directory = temp_dir();
214  
215          let program_string = r"
216  program token.alpha;
217  
218  record token:
219      owner as address.private;
220      token_amount as u64.private;
221  
222  function compute:
223      input r0 as token.record;
224      add r0.token_amount r0.token_amount into r1;
225      output r1 as u64.private;";
226  
227          // Initialize a new program.
228          let (string, program) = Program::<CurrentNetwork>::parse(program_string).unwrap();
229          assert!(string.is_empty(), "Parser did not consume all of the string: '{string}'");
230  
231          // Construct the process.
232          let mut process = Process::load().unwrap();
233          // Add the program to the process.
234          process.add_program(&program).unwrap();
235  
236          // Prepare the function name.
237          let function_name = Identifier::from_str("compute").unwrap();
238  
239          // Sample the proving key.
240          process.synthesize_key::<CurrentAleo, _>(program.id(), &function_name, &mut TestRng::default()).unwrap();
241  
242          // Retrieve the proving key.
243          let proving_key = process.get_proving_key(program.id(), function_name).unwrap();
244  
245          // Create the prover file at the path.
246          let expected = ProverFile::create(&directory, &function_name, proving_key).unwrap();
247          // Open the prover file at the path.
248          let candidate = ProverFile::open(&directory, &function_name).unwrap();
249          // Ensure the prover files are equal.
250          assert_eq!(expected.to_bytes_le().unwrap(), candidate.to_bytes_le().unwrap());
251      }
252  }