refactor(contracts): improve contracts and test documentation; small code optimization and refactor

Author: 0xjei

packages/contracts/contracts/src/AdvancedChecker.sol

@@ -3,41 +3,45 @@ pragma solidity 0.8.27;
 
 import {IAdvancedChecker, Check} from "./interfaces/IAdvancedChecker.sol";
 
+/// @notice Tracks validation status for pre, main, and post checks.
+/// @dev Used to maintain check state in AdvancedPolicy.
 struct CheckStatus {
+    /// @dev Pre-check completion status.
     bool pre;
+    /// @dev Number of completed main checks.
     uint8 main;
+    /// @dev Post-check completion status.
     bool post;
 }
 
 /// @title AdvancedChecker.
-/// @notice Abstract base contract which can be extended to implement a specific `AdvancedChecker`.
-/// @dev The `AdvancedChecker` contract builds upon the `BaseChecker` by introducing additional validation phases.
-/// It allows for pre-condition (`PRE`), main (`MAIN`), and post-condition (`POST`) checks, with the option to skip
-/// pre and post checks based on constructor parameters. The `_check` method orchestrates the validation process
-/// based on the specified check type.
+/// @notice Multi-phase validation checker with pre, main, and post checks.
+/// @dev Base contract for implementing complex validation logic with configurable check phases.
 abstract contract AdvancedChecker is IAdvancedChecker {
-    /// @notice Flag to determine if pre-condition checks should be skipped.
+    /// @notice Controls whether pre-condition checks are required.
     bool public immutable SKIP_PRE;
 
-    /// @notice Flag to determine if post-condition checks should be skipped.
+    /// @notice Controls whether post-condition checks are required.
     bool public immutable SKIP_POST;
 
-    /// @notice Flag to determine if main checks can be executed multiple times.
+    /// @notice Controls whether main check can be executed multiple times.
     bool public immutable ALLOW_MULTIPLE_MAIN;
 
-    /// @param _skipPre Indicates whether to skip pre-condition checks.
-    /// @param _skipPost Indicates whether to skip post-condition checks.
-    /// @param _allowMultipleMain Indicates whether the main check can be executed multiple times.
+    /// @notice Sets up checker configuration.
+    /// @param _skipPre Skip pre-condition validation.
+    /// @param _skipPost Skip post-condition validation.
+    /// @param _allowMultipleMain Allow multiple main validations.
     constructor(bool _skipPre, bool _skipPost, bool _allowMultipleMain) {
         SKIP_PRE = _skipPre;
         SKIP_POST = _skipPost;
         ALLOW_MULTIPLE_MAIN = _allowMultipleMain;
     }
 
-    /// @notice Public method to check the validity of the provided evidence for a given address and check type.
-    /// @param subject The address to be checked.
-    /// @param evidence The evidence associated with the check.
-    /// @param checkType The type of check to perform (PRE, MAIN, POST).
+    /// @notice Entry point for validation checks.
+    /// @param subject Address to validate.
+    /// @param evidence Validation data.
+    /// @param checkType Type of check (PRE, MAIN, POST).
+    /// @return checked Validation result.
     function check(
         address subject,
         bytes memory evidence,
@@ -46,37 +50,46 @@ abstract contract AdvancedChecker is IAdvancedChecker {
         return _check(subject, evidence, checkType);
     }
 
-    /// @notice Internal method to orchestrate the validation process based on the specified check type.
-    /// @param subject The address to be checked.
-    /// @param evidence The evidence associated with the check.
-    /// @param checkType The type of check to perform (PRE, MAIN, POST).
+    /// @notice Core validation logic router.
+    /// @dev Directs to appropriate check based on type and configuration.
+    /// @param subject Address to validate.
+    /// @param evidence Validation data.
+    /// @param checkType Check type to perform.
+    /// @return checked Validation result.
+    /// @custom:throws PreCheckSkipped If PRE check attempted when skipped.
+    /// @custom:throws PostCheckSkipped If POST check attempted when skipped.
     function _check(address subject, bytes memory evidence, Check checkType) internal view returns (bool checked) {
-        if (SKIP_PRE && checkType == Check.PRE) revert PreCheckSkipped();
-        if (SKIP_POST && checkType == Check.POST) revert PostCheckSkipped();
+        // Validate skip conditions first.
+        if (checkType == Check.PRE && SKIP_PRE) revert PreCheckSkipped();
+        if (checkType == Check.POST && SKIP_POST) revert PostCheckSkipped();
 
-        if (!SKIP_PRE && checkType == Check.PRE) {
-            return _checkPre(subject, evidence);
-        }
-
-        if (!SKIP_POST && checkType == Check.POST) {
-            return _checkPost(subject, evidence);
-        }
-
-        return _checkMain(subject, evidence);
+        // Route to appropriate check.
+        return
+            checkType == Check.PRE
+                ? _checkPre(subject, evidence)
+                : checkType == Check.POST
+                    ? _checkPost(subject, evidence)
+                    : _checkMain(subject, evidence);
     }
 
-    /// @notice Internal method for performing pre-condition checks.
-    /// @param subject The address to be checked.
-    /// @param evidence The evidence associated with the check.
+    /// @notice Pre-condition validation implementation.
+    /// @dev Override to implement pre-check logic.
+    /// @param subject Address to validate.
+    /// @param evidence Validation data.
+    /// @return checked Validation result.
     function _checkPre(address subject, bytes memory evidence) internal view virtual returns (bool checked) {}
 
-    /// @notice Internal method for performing main checks.
-    /// @param subject The address to be checked.
-    /// @param evidence The evidence associated with the check.
+    /// @notice Main validation implementation.
+    /// @dev Override to implement main check logic.
+    /// @param subject Address to validate.
+    /// @param evidence Validation data.
+    /// @return checked Validation result.
     function _checkMain(address subject, bytes memory evidence) internal view virtual returns (bool checked) {}
 
-    /// @notice Internal method for performing post-condition checks.
-    /// @param subject The address to be checked.
-    /// @param evidence The evidence associated with the check.
+    /// @notice Post-condition validation implementation.
+    /// @dev Override to implement post-check logic.
+    /// @param subject Address to validate.
+    /// @param evidence Validation data.
+    /// @return checked Validation result.
     function _checkPost(address subject, bytes memory evidence) internal view virtual returns (bool checked) {}
 }

packages/contracts/contracts/src/AdvancedPolicy.sol

@@ -5,80 +5,84 @@ import {Policy} from "./Policy.sol";
 import {IAdvancedPolicy, Check} from "./interfaces/IAdvancedPolicy.sol";
 import {AdvancedChecker, CheckStatus} from "./AdvancedChecker.sol";
 
-/// @title AdvancedPolicy
-/// @notice Abstract base contract which can be extended to implement a specific `AdvancedPolicy`.
+/// @title AdvancedPolicy.
+/// @notice Implements advanced policy checks with pre, main, and post validation stages.
+/// @dev Extends Policy contract with multi-stage validation capabilities.
 abstract contract AdvancedPolicy is IAdvancedPolicy, Policy {
-    /// @dev Reference to the AdvancedChecker contract for validation.
+    /// @notice Reference to the validation checker contract.
+    /// @dev Immutable to ensure checker cannot be changed after deployment.
     AdvancedChecker public immutable ADVANCED_CHECKER;
 
-    /// @dev Tracks the check status of each address.
+    /// @notice Tracks validation status for each subject per target.
+    /// @dev Maps target => subject => CheckStatus.
     mapping(address => mapping(address => CheckStatus)) public enforced;
 
-    /// @notice Constructor to initialize the AdvancedChecker contract.
-    /// @param _advancedChecker The address of the AdvancedChecker contract.
+    /// @notice Initializes contract with an AdvancedChecker instance.
+    /// @param _advancedChecker Address of the AdvancedChecker contract.
     constructor(AdvancedChecker _advancedChecker) {
         ADVANCED_CHECKER = _advancedChecker;
     }
 
-    /// @notice Enforces the custom target logic.
-    /// @dev Calls the internal `_enforce` function to enforce the target logic.
-    /// @dev Must call the `check` to handle the logic of checking subject for specific target.
-    /// @param subject The address of those who have successfully enforced the check.
-    /// @param evidence Additional data required for the check (e.g., encoded token identifier).
-    /// @param checkType The type of the check to be enforced for the subject with the given data.
+    /// @notice Enforces policy check for a subject.
+    /// @dev Only callable by target contract.
+    /// @param subject Address to validate.
+    /// @param evidence Validation data.
+    /// @param checkType Type of check (PRE, MAIN, POST).
     function enforce(address subject, bytes calldata evidence, Check checkType) external override onlyTarget {
         _enforce(subject, evidence, checkType);
     }
 
-    /// @notice Internal function to enforce the target logic.
-    /// @param subject The address of those who have successfully enforced the check.
-    /// @param evidence Additional data required for the check (e.g., encoded token identifier).
-    /// @param checkType The type of the check to be enforced for the subject with the given data.
+    /// @notice Internal check enforcement logic.
+    /// @dev Handles different check types and their dependencies.
+    /// @param subject Address to validate.
+    /// @param evidence Validation data.
+    /// @param checkType Type of check to perform.
+    /// @custom:throws UnsuccessfulCheck If validation fails.
+    /// @custom:throws AlreadyEnforced If check was already completed.
+    /// @custom:throws PreCheckNotEnforced If PRE check is required but not done.
+    /// @custom:throws MainCheckNotEnforced If MAIN check is required but not done.
+    /// @custom:throws MainCheckAlreadyEnforced If multiple MAIN checks not allowed.
     function _enforce(address subject, bytes calldata evidence, Check checkType) internal {
-        bool checked = ADVANCED_CHECKER.check(subject, evidence, checkType);
-
-        if (!checked) {
+        if (!ADVANCED_CHECKER.check(subject, evidence, checkType)) {
             revert UnsuccessfulCheck();
         }
 
+        CheckStatus storage status = enforced[msg.sender][subject];
+
+        // Handle PRE check.
         if (checkType == Check.PRE) {
-            if (!ADVANCED_CHECKER.SKIP_POST() && enforced[msg.sender][subject].pre) {
+            if (!ADVANCED_CHECKER.SKIP_POST() && status.pre) {
                 revert AlreadyEnforced();
-            } else {
-                enforced[msg.sender][subject].pre = true;
             }
-        } else {
-            if (checkType == Check.POST) {
-                if (enforced[msg.sender][subject].post) {
-                    revert AlreadyEnforced();
-                } else {
-                    if (!ADVANCED_CHECKER.SKIP_PRE() && !enforced[msg.sender][subject].pre) {
-                        revert PreCheckNotEnforced();
-                    } else {
-                        if (enforced[msg.sender][subject].main == 0) {
-                            revert MainCheckNotEnforced();
-                        } else {
-                            enforced[msg.sender][subject].post = true;
-                        }
-                    }
-                }
-            } else {
-                if (
-                    checkType == Check.MAIN &&
-                    !ADVANCED_CHECKER.ALLOW_MULTIPLE_MAIN() &&
-                    enforced[msg.sender][subject].main > 0
-                ) {
-                    revert MainCheckAlreadyEnforced();
-                } else {
-                    if (checkType == Check.MAIN && !ADVANCED_CHECKER.SKIP_PRE() && !enforced[msg.sender][subject].pre) {
-                        revert PreCheckNotEnforced();
-                    } else {
-                        enforced[msg.sender][subject].main += 1;
-                    }
-                }
+            status.pre = true;
+            emit Enforced(subject, target, evidence, checkType);
+            return;
+        }
+
+        // Handle POST check.
+        if (checkType == Check.POST) {
+            if (status.post) {
+                revert AlreadyEnforced();
             }
+            if (!ADVANCED_CHECKER.SKIP_PRE() && !status.pre) {
+                revert PreCheckNotEnforced();
+            }
+            if (status.main == 0) {
+                revert MainCheckNotEnforced();
+            }
+            status.post = true;
+            emit Enforced(subject, target, evidence, checkType);
+            return;
         }
 
+        // Handle MAIN check.
+        if (!ADVANCED_CHECKER.ALLOW_MULTIPLE_MAIN() && status.main > 0) {
+            revert MainCheckAlreadyEnforced();
+        }
+        if (!ADVANCED_CHECKER.SKIP_PRE() && !status.pre) {
+            revert PreCheckNotEnforced();
+        }
+        status.main += 1;
         emit Enforced(subject, target, evidence, checkType);
     }
 }

packages/contracts/contracts/src/BaseChecker.sol

@@ -3,21 +3,24 @@ pragma solidity 0.8.27;
 
 import {IBaseChecker} from "./interfaces/IBaseChecker.sol";
 
-/// @title BaseChecker.
-/// @notice Abstract base contract which can be extended to implement a specific `BaseChecker`.
-/// @dev The `BaseChecker` contract provides a foundational structure for implementing specific checker logic.
-/// It defines a method `check` that invokes a protected `_check` method, which must be implemented by derived
-/// contracts.
+/// @title BaseChecker
+/// @notice Abstract base contract for implementing validation checks.
+/// @dev Provides a standardized interface for implementing custom validation logic
+/// through the internal _check method.
 abstract contract BaseChecker is IBaseChecker {
-    /// @notice Checks the validity of the provided evidence for a given address.
-    /// @param subject The address to be checked.
-    /// @param evidence The evidence associated with the check.
+    /// @notice Validates evidence for a given subject address.
+    /// @dev External view function that delegates to internal _check implementation.
+    /// @param subject Address to validate.
+    /// @param evidence Custom validation data.
+    /// @return checked Boolean indicating if the check passed.
     function check(address subject, bytes memory evidence) external view override returns (bool checked) {
         return _check(subject, evidence);
     }
 
-    /// @notice Internal method to perform the actual check logic.
-    /// @param subject The address to be checked.
-    /// @param evidence The evidence associated with the check.
+    /// @notice Internal validation logic implementation.
+    /// @dev Must be implemented by derived contracts.
+    /// @param subject Address to validate.
+    /// @param evidence Custom validation data.
+    /// @return checked Boolean indicating if the check passed.
     function _check(address subject, bytes memory evidence) internal view virtual returns (bool checked) {}
 }

packages/contracts/contracts/src/BasePolicy.sol

@@ -6,31 +6,45 @@ import {Policy} from "./Policy.sol";
 import {BaseChecker} from "./BaseChecker.sol";
 
 /// @title BasePolicy
-/// @notice Abstract base contract which can be extended to implement a specific `BasePolicy`.
+/// @notice Abstract base contract for implementing specific policy checks.
+/// @dev Inherits from Policy and implements IBasePolicy interface.
+///
+/// Provides core functionality for enforcing policy checks through a BaseChecker
+/// contract. Each specific policy implementation should extend this contract
+/// and implement its custom checking logic.
 abstract contract BasePolicy is Policy, IBasePolicy {
-    /// @dev Reference to the BaseChecker contract for validation.
+    /// @notice Reference to the BaseChecker contract used for validation.
+    /// @dev Immutable to ensure checker cannot be changed after deployment.
     BaseChecker public immutable BASE_CHECKER;
 
-    /// @dev Tracks whether the check has been enforced for a subject.
+    /// @notice Tracks enforcement status for each subject per target.
+    /// @dev Maps target => subject => enforcement status.
     mapping(address => mapping(address => bool)) public enforced;
 
-    /// @notice Constructor to initialize the BaseChecker contract.
-    /// @param _baseChecker The address of the BaseChecker contract.
+    /// @notice Initializes the contract with a BaseChecker instance.
+    /// @param _baseChecker Address of the BaseChecker contract.
+    /// @dev The BaseChecker address cannot be changed after deployment.
     constructor(BaseChecker _baseChecker) {
         BASE_CHECKER = _baseChecker;
     }
 
-    /// @notice Enforces the custom target enforcing logic.
-    /// @dev Must call the `check` to handle the logic of checking subject for specific target.
-    /// @param subject The address of those who have successfully enforced the check.
-    /// @param evidence Additional data required for the check (e.g., encoded token identifier).
+    /// @notice External function to enforce policy checks.
+    /// @dev Only callable by the target contract.
+    /// @param subject Address to enforce the check on.
+    /// @param evidence Additional data required for verification.
+    /// @custom:throws AlreadyEnforced if check was previously enforced.
+    /// @custom:throws UnsuccessfulCheck if the check fails.
+    /// @custom:emits Enforced when check succeeds.
     function enforce(address subject, bytes calldata evidence) external override onlyTarget {
         _enforce(subject, evidence);
     }
 
-    /// @notice Enforces the custom target enforcing logic.
-    /// @param subject The address of those who have successfully enforced the check.
-    /// @param evidence Additional data required for the check (e.g., encoded token identifier).
+    /// @notice Internal implementation of enforcement logic.
+    /// @dev Performs the actual check using BASE_CHECKER.
+    /// @param subject Address to enforce the check on.
+    /// @param evidence Additional data required for verification.
+    /// @custom:throws AlreadyEnforced if already enforced for this subject.
+    /// @custom:throws UnsuccessfulCheck if BASE_CHECKER.check returns false.
     function _enforce(address subject, bytes calldata evidence) internal {
         bool checked = BASE_CHECKER.check(subject, evidence);