Repository: ergoplatform/sigma-rust (language: rust)

Issue:
Move `Input`, `DataInput` to `transaction` module and `ProverResult` to `sigma_protocol` module



Discussion / hints:

Fix patch:
diff --git a/bindings/ergo-lib-wasm/src/data_input.rs b/bindings/ergo-lib-wasm/src/data_input.rs
index 330793b8a..6ab1be1ac 100644
--- a/bindings/ergo-lib-wasm/src/data_input.rs
+++ b/bindings/ergo-lib-wasm/src/data_input.rs
@@ -6,7 +6,7 @@ use wasm_bindgen::prelude::*;
 /// Inputs, that are used to enrich script context, but won't be spent by the transaction
 #[wasm_bindgen]
 #[derive(PartialEq, Eq, Debug, Clone)]
-pub struct DataInput(chain::data_input::DataInput);
+pub struct DataInput(chain::transaction::DataInput);
 
 /// DataInput collection
 #[wasm_bindgen]
@@ -36,7 +36,7 @@ impl DataInputs {
     }
 }
 
-impl From<&DataInputs> for Vec<chain::data_input::DataInput> {
+impl From<&DataInputs> for Vec<chain::transaction::DataInput> {
     fn from(v: &DataInputs) -> Self {
         v.0.clone().iter().map(|i| i.0.clone()).collect()
     }
diff --git a/bindings/ergo-lib-wasm/src/ergo_tree.rs b/bindings/ergo-lib-wasm/src/ergo_tree.rs
index 8608e3a3c..d9c566cdd 100644
--- a/bindings/ergo-lib-wasm/src/ergo_tree.rs
+++ b/bindings/ergo-lib-wasm/src/ergo_tree.rs
@@ -12,7 +12,7 @@ use derive_more::{From, Into};
 /// The root of ErgoScript IR. Serialized instances of this class are self sufficient and can be passed around.
 #[wasm_bindgen]
 #[derive(PartialEq, Eq, Debug, Clone, From, Into)]
-pub struct ErgoTree(ergo_lib::ErgoTree);
+pub struct ErgoTree(ergo_lib::ergo_tree::ErgoTree);
 
 #[wasm_bindgen]
 impl ErgoTree {
@@ -20,7 +20,7 @@ impl ErgoTree {
     pub fn from_base16_bytes(s: &str) -> Result<ErgoTree, JsValue> {
         let bytes = Base16DecodedBytes::try_from(s.to_string())
             .map_err(|e| JsValue::from_str(&format!("{}", e)))?;
-        ergo_lib::ErgoTree::sigma_parse_bytes(bytes.0)
+        ergo_lib::ergo_tree::ErgoTree::sigma_parse_bytes(bytes.0)
             .map(ErgoTree)
             .map_err(|e| JsValue::from_str(&format!("{}", e)))
     }
diff --git a/bindings/ergo-lib-wasm/src/secret_key.rs b/bindings/ergo-lib-wasm/src/secret_key.rs
index c47bd93d5..da7c5a1d0 100644
--- a/bindings/ergo-lib-wasm/src/secret_key.rs
+++ b/bindings/ergo-lib-wasm/src/secret_key.rs
@@ -1,7 +1,7 @@
 //! Secret key
 use std::convert::TryInto;
 
-use ergo_lib::sigma_protocol::DlogProverInput;
+use ergo_lib::sigma_protocol::private_input::DlogProverInput;
 use ergo_lib::wallet;
 use wasm_bindgen::prelude::*;
 
diff --git a/ergo-lib/CHANGELOG.md b/ergo-lib/CHANGELOG.md
index 6e13e6215..5cd0794d6 100644
--- a/ergo-lib/CHANGELOG.md
+++ b/ergo-lib/CHANGELOG.md
@@ -8,11 +8,16 @@ and this project adheres to [Semantic Versioning](http://semver.org/).
 ## [Unreleased] - ReleaseDate
 
 ### Added
+
 - Support for parsing ErgoBox id also from "id" JSON field name [#134](https://github.com/ergoplatform/sigma-rust/pull/134)
 - Address::p2pk_from_ergo_tree to make an Address from ErgoTree with root ProveDlog [#136](https://github.com/ergoplatform/sigma-rust/pull/136)
 - Address::p2pk_from_pk_bytes to make an Address from serialized PK [#136](https://github.com/ergoplatform/sigma-rust/pull/136)
 - Address::from_str to parse an Address without checking the network prefix [#136](https://github.com/ergoplatform/sigma-rust/pull/136)
 
+### Changed
+
+- moved and change visibility of various modules(input, data_input, prover_result, etc.) [#135](https://github.com/ergoplatform/sigma-rust/pull/135) 
+
 ## [0.3.0] - 2020-11-04
 
 ### Added
diff --git a/ergo-lib/src/chain.rs b/ergo-lib/src/chain.rs
index bf43f07ac..55aa38f15 100644
--- a/ergo-lib/src/chain.rs
+++ b/ergo-lib/src/chain.rs
@@ -11,12 +11,8 @@ pub use base16_bytes::Base16EncodedBytes;
 pub use digest32::*;
 
 pub mod address;
-pub mod context_extension;
 pub mod contract;
-pub mod data_input;
 pub mod ergo_box;
 pub mod ergo_state_context;
-pub mod input;
-pub mod prover_result;
 pub mod token;
 pub mod transaction;
diff --git a/ergo-lib/src/chain/address.rs b/ergo-lib/src/chain/address.rs
index b8f785339..25aa18e70 100644
--- a/ergo-lib/src/chain/address.rs
+++ b/ergo-lib/src/chain/address.rs
@@ -1,17 +1,14 @@
 //! Address types
 
 use super::digest32;
-use crate::ast::Constant;
-use crate::ast::ConstantVal;
-use crate::ErgoTreeParsingError;
 use crate::{
-    ast::Expr,
+    ast::{Constant, ConstantVal, Expr},
+    ergo_tree::{ErgoTree, ErgoTreeParsingError},
     serialization::{SerializationError, SigmaSerializable},
     sigma_protocol::{
         dlog_group::EcPoint,
         sigma_boolean::{ProveDlog, SigmaBoolean, SigmaProofOfKnowledgeTree, SigmaProp},
     },
-    ErgoTree,
 };
 use std::{
     convert::{TryFrom, TryInto},
diff --git a/ergo-lib/src/chain/contract.rs b/ergo-lib/src/chain/contract.rs
index 231b5fd11..4df9ff9fd 100644
--- a/ergo-lib/src/chain/contract.rs
+++ b/ergo-lib/src/chain/contract.rs
@@ -1,7 +1,8 @@
 //! Ergo contract
 
 use super::address::Address;
-use crate::{serialization::SerializationError, ErgoTree};
+use crate::ergo_tree::ErgoTree;
+use crate::serialization::SerializationError;
 
 /// High-level wrapper for ErgoTree
 #[derive(PartialEq, Eq, Debug, Clone)]
diff --git a/ergo-lib/src/chain/ergo_box/box_builder.rs b/ergo-lib/src/chain/ergo_box/box_builder.rs
index 4362e459e..0bc14ec58 100644
--- a/ergo-lib/src/chain/ergo_box/box_builder.rs
+++ b/ergo-lib/src/chain/ergo_box/box_builder.rs
@@ -1,18 +1,15 @@
 //! ErgoBoxCandidate builder
 
 use std::collections::HashMap;
-use std::convert::TryFrom;
-use std::convert::TryInto;
+use std::convert::{TryFrom, TryInto};
 
 use crate::ast::Constant;
 use crate::chain::token::Token;
+use crate::ergo_tree::ErgoTree;
 use crate::serialization::SigmaSerializable;
-use crate::ErgoTree;
 
 use super::box_value::BoxValue;
-use super::register::NonMandatoryRegisterId;
-use super::register::NonMandatoryRegisters;
-use super::register::NonMandatoryRegistersError;
+use super::register::{NonMandatoryRegisterId, NonMandatoryRegisters, NonMandatoryRegistersError};
 use super::ErgoBoxCandidate;
 use thiserror::Error;
 
diff --git a/ergo-lib/src/chain/json.rs b/ergo-lib/src/chain/json.rs
index 0f6e33f2d..7a19acdf2 100644
--- a/ergo-lib/src/chain/json.rs
+++ b/ergo-lib/src/chain/json.rs
@@ -13,7 +13,8 @@ where
 pub mod ergo_tree {
 
     use super::*;
-    use crate::{serialization::SigmaSerializable, ErgoTree};
+    use crate::ergo_tree::ErgoTree;
+    use crate::serialization::SigmaSerializable;
     use serde::{Deserialize, Deserializer, Serializer};
 
     pub fn serialize<S>(ergo_tree: &ErgoTree, serializer: S) -> Result<S::Ok, S::Error>
@@ -38,12 +39,13 @@ pub mod ergo_tree {
 }
 
 pub mod ergo_box {
-    use crate::chain::ergo_box::BoxId;
-    use crate::chain::ergo_box::BoxValue;
-    use crate::chain::ergo_box::NonMandatoryRegisters;
     use crate::{
-        chain::{token::Token, transaction::TxId},
-        ErgoTree,
+        chain::{
+            ergo_box::{BoxId, BoxValue, NonMandatoryRegisters},
+            token::Token,
+            transaction::TxId,
+        },
+        ergo_tree::ErgoTree,
     };
     use serde::Deserialize;
 
@@ -78,8 +80,8 @@ pub mod ergo_box {
 }
 
 pub mod transaction {
-    use crate::chain::input::UnsignedInput;
-    use crate::chain::{data_input::DataInput, ergo_box::ErgoBox, input::Input, transaction::TxId};
+    use crate::chain::transaction::{DataInput, Input, UnsignedInput};
+    use crate::chain::{ergo_box::ErgoBox, transaction::TxId};
     use serde::{Deserialize, Serialize};
 
     #[derive(Serialize, Deserialize, PartialEq, Debug, Clone)]
@@ -120,12 +122,12 @@ pub mod transaction {
 #[cfg(test)]
 mod tests {
     use crate::chain::transaction::unsigned::UnsignedTransaction;
+    use crate::sigma_protocol::prover::ContextExtension;
     use std::convert::TryInto;
 
     use super::super::ergo_box::*;
     use super::super::transaction::*;
     use super::*;
-    use crate::chain::context_extension::ContextExtension;
     use proptest::prelude::*;
 
     proptest! {
diff --git a/ergo-lib/src/chain/transaction.rs b/ergo-lib/src/chain/transaction.rs
index 533924ffc..196843164 100644
--- a/ergo-lib/src/chain/transaction.rs
+++ b/ergo-lib/src/chain/transaction.rs
@@ -1,15 +1,18 @@
 //! Ergo transaction
 
+mod data_input;
+mod input;
 pub mod unsigned;
 
+pub use data_input::*;
+pub use input::*;
+
 #[cfg(feature = "json")]
 use super::json;
 use super::{
-    data_input::DataInput,
     digest32::{blake2b256_hash, Digest32},
     ergo_box::ErgoBox,
     ergo_box::ErgoBoxCandidate,
-    input::Input,
     token::TokenId,
 };
 use crate::serialization::{
diff --git a/ergo-lib/src/chain/data_input.rs b/ergo-lib/src/chain/transaction/data_input.rs
similarity index 97%
rename from ergo-lib/src/chain/data_input.rs
rename to ergo-lib/src/chain/transaction/data_input.rs
index 3c2e4be9f..f0bcc48f8 100644
--- a/ergo-lib/src/chain/data_input.rs
+++ b/ergo-lib/src/chain/transaction/data_input.rs
@@ -2,7 +2,7 @@
 
 use std::io;
 
-use super::ergo_box::BoxId;
+use crate::chain::ergo_box::BoxId;
 use crate::serialization::{
     sigma_byte_reader::SigmaByteRead, sigma_byte_writer::SigmaByteWrite, SerializationError,
     SigmaSerializable,
diff --git a/ergo-lib/src/chain/input.rs b/ergo-lib/src/chain/transaction/input.rs
similarity index 92%
rename from ergo-lib/src/chain/input.rs
rename to ergo-lib/src/chain/transaction/input.rs
index bbdca69de..c198d22f3 100644
--- a/ergo-lib/src/chain/input.rs
+++ b/ergo-lib/src/chain/transaction/input.rs
@@ -1,16 +1,14 @@
 //! Transaction input
 use std::io;
 
-use super::{
-    context_extension::ContextExtension,
-    ergo_box::BoxId,
-    ergo_box::ErgoBoxId,
-    prover_result::{ProofBytes, ProverResult},
-};
+use crate::chain::ergo_box::{BoxId, ErgoBoxId};
 use crate::serialization::{
     sigma_byte_reader::SigmaByteRead, sigma_byte_writer::SigmaByteWrite, SerializationError,
     SigmaSerializable,
 };
+use crate::sigma_protocol::prover::ContextExtension;
+use crate::sigma_protocol::prover::ProofBytes;
+use crate::sigma_protocol::prover::ProverResult;
 #[cfg(feature = "json")]
 use serde::{Deserialize, Serialize};
 
diff --git a/ergo-lib/src/chain/transaction/unsigned.rs b/ergo-lib/src/chain/transaction/unsigned.rs
index 08f62af5d..3d3f18f2a 100644
--- a/ergo-lib/src/chain/transaction/unsigned.rs
+++ b/ergo-lib/src/chain/transaction/unsigned.rs
@@ -1,24 +1,19 @@
 //! Unsigned (without proofs) transaction
 
+use super::input::{Input, UnsignedInput};
 #[cfg(feature = "json")]
 use super::json;
+use super::DataInput;
 use super::{
-    super::{
-        data_input::DataInput, digest32::blake2b256_hash, ergo_box::ErgoBoxCandidate, input::Input,
-    },
+    super::{digest32::blake2b256_hash, ergo_box::ErgoBoxCandidate},
     Transaction, TxId,
 };
 #[cfg(feature = "json")]
 use crate::chain::transaction::ErgoBox;
 #[cfg(feature = "json")]
 use crate::chain::transaction::TransactionFromJsonError;
-use crate::{
-    chain::{
-        input::UnsignedInput,
-        prover_result::{ProofBytes, ProverResult},
-    },
-    serialization::SigmaSerializable,
-};
+use crate::serialization::SigmaSerializable;
+use crate::sigma_protocol::prover::{ProofBytes, ProverResult};
 #[cfg(feature = "json")]
 use core::convert::TryFrom;
 #[cfg(feature = "json")]
diff --git a/ergo-lib/src/lib.rs b/ergo-lib/src/lib.rs
index 35668e284..1cbd46855 100644
--- a/ergo-lib/src/lib.rs
+++ b/ergo-lib/src/lib.rs
@@ -14,17 +14,17 @@
 
 mod big_integer;
 mod constants;
-mod ergo_tree;
-pub mod serialization;
 mod types;
 
 pub mod ast;
 pub mod chain;
 pub mod eval;
+pub mod serialization;
 pub mod sigma_protocol;
 pub mod util;
 pub mod wallet;
-pub use ergo_tree::*;
+
+pub mod ergo_tree;
 
 #[cfg(test)]
 pub mod test_util;
diff --git a/ergo-lib/src/serialization.rs b/ergo-lib/src/serialization.rs
index 0e40a9264..01d621ba3 100644
--- a/ergo-lib/src/serialization.rs
+++ b/ergo-lib/src/serialization.rs
@@ -10,9 +10,9 @@ mod sigmaboolean;
 pub mod constant_store;
 pub mod ergo_box;
 pub mod op_code;
-pub mod serializable;
 pub mod sigma_byte_reader;
 pub mod sigma_byte_writer;
 pub mod types;
 
+mod serializable;
 pub use serializable::*;
diff --git a/ergo-lib/src/serialization/ergo_box.rs b/ergo-lib/src/serialization/ergo_box.rs
index 51aa3e8b1..b856d6a50 100644
--- a/ergo-lib/src/serialization/ergo_box.rs
+++ b/ergo-lib/src/serialization/ergo_box.rs
@@ -1,20 +1,17 @@
 //! ErgoBox and ErgoBoxCandidate serialization
-use crate::serialization::{
-    sigma_byte_reader::SigmaByteRead, SerializationError, SigmaSerializable,
-};
 use crate::{
     ast::Constant,
     chain::{
         ergo_box::{BoxValue, ErgoBoxCandidate, NonMandatoryRegisters},
         token::{Token, TokenId},
     },
-    ErgoTree,
+    ergo_tree::ErgoTree,
+    serialization::{sigma_byte_reader::SigmaByteRead, SerializationError, SigmaSerializable},
 };
 use indexmap::IndexSet;
 
 use super::sigma_byte_writer::SigmaByteWrite;
-use std::convert::TryFrom;
-use std::convert::TryInto;
+use std::convert::{TryFrom, TryInto};
 use std::io;
 
 /// Box serialization with token ids optionally saved in transaction
diff --git a/ergo-lib/src/sigma_protocol.rs b/ergo-lib/src/sigma_protocol.rs
index 65d1624e1..88bb2faa3 100644
--- a/ergo-lib/src/sigma_protocol.rs
+++ b/ergo-lib/src/sigma_protocol.rs
@@ -1,21 +1,18 @@
 //! Sigma protocols
 
-mod challenge;
-mod private_input;
-
-pub use challenge::*;
-pub use private_input::*;
-
 pub mod dlog_group;
-pub mod dlog_protocol;
-pub mod fiat_shamir;
+pub mod private_input;
 pub mod prover;
-pub mod sig_serializer;
 pub mod sigma_boolean;
-pub mod unchecked_tree;
-pub mod unproven_tree;
 pub mod verifier;
 
+mod challenge;
+mod dlog_protocol;
+mod fiat_shamir;
+mod sig_serializer;
+mod unchecked_tree;
+mod unproven_tree;
+
 use k256::Scalar;
 
 use dlog_protocol::FirstDlogProverMessage;
@@ -24,6 +21,8 @@ use std::convert::TryInto;
 use unchecked_tree::{UncheckedSigmaTree, UncheckedTree};
 use unproven_tree::{UnprovenLeaf, UnprovenSchnorr, UnprovenTree};
 
+use self::challenge::Challenge;
+
 /** The message sent by a prover to its associated verifier as part of a sigma protocol interaction. */
 pub trait ProverMessage {
     /// serialized message
diff --git a/ergo-lib/src/sigma_protocol/dlog_group.rs b/ergo-lib/src/sigma_protocol/dlog_group.rs
index ffb5202f3..e5d8565ac 100644
--- a/ergo-lib/src/sigma_protocol/dlog_group.rs
+++ b/ergo-lib/src/sigma_protocol/dlog_group.rs
@@ -15,7 +15,6 @@
 //!
 //! On the other hand, any group element can be mapped to some string.
 
-use super::DlogProverInput;
 use crate::serialization::{
     sigma_byte_reader::SigmaByteRead, SerializationError, SigmaSerializable,
 };
@@ -28,6 +27,8 @@ use std::{
     ops::{Add, Mul, Neg},
 };
 
+use super::private_input::DlogProverInput;
+
 /// Elliptic curve point
 #[derive(PartialEq, Debug, Clone)]
 pub struct EcPoint(ProjectivePoint);
diff --git a/ergo-lib/src/sigma_protocol/dlog_protocol.rs b/ergo-lib/src/sigma_protocol/dlog_protocol.rs
index 80a55a76e..9389aa033 100644
--- a/ergo-lib/src/sigma_protocol/dlog_protocol.rs
+++ b/ergo-lib/src/sigma_protocol/dlog_protocol.rs
@@ -42,7 +42,7 @@ impl From<Scalar> for SecondDlogProverMessage {
 /// Interactive prover
 pub mod interactive_prover {
     use super::{FirstDlogProverMessage, SecondDlogProverMessage};
-    use crate::sigma_protocol::{dlog_group, Challenge, DlogProverInput, ProveDlog};
+    use crate::sigma_protocol::{dlog_group, private_input::DlogProverInput, Challenge, ProveDlog};
     use dlog_group::EcPoint;
     use k256::Scalar;
 
@@ -101,7 +101,7 @@ pub mod interactive_prover {
 mod tests {
     use super::super::*;
     use super::*;
-    use crate::sigma_protocol::DlogProverInput;
+    use crate::sigma_protocol::private_input::DlogProverInput;
     use proptest::prelude::*;
 
     proptest! {
diff --git a/ergo-lib/src/sigma_protocol/fiat_shamir.rs b/ergo-lib/src/sigma_protocol/fiat_shamir.rs
index 4d389b8b2..3f7db902b 100644
--- a/ergo-lib/src/sigma_protocol/fiat_shamir.rs
+++ b/ergo-lib/src/sigma_protocol/fiat_shamir.rs
@@ -6,7 +6,7 @@ use super::{
     unproven_tree::UnprovenTree,
     ProofTree, ProofTreeLeaf, ProverMessage, GROUP_SIZE, SOUNDNESS_BYTES,
 };
-use crate::{ast::Expr, serialization::SigmaSerializable, ErgoTree};
+use crate::{ast::Expr, ergo_tree::ErgoTree, serialization::SigmaSerializable};
 use blake2::digest::{Update, VariableOutput};
 use blake2::VarBlake2b;
 use std::{
diff --git a/ergo-lib/src/sigma_protocol/private_input.rs b/ergo-lib/src/sigma_protocol/private_input.rs
index b9a224f5d..0d179afca 100644
--- a/ergo-lib/src/sigma_protocol/private_input.rs
+++ b/ergo-lib/src/sigma_protocol/private_input.rs
@@ -1,6 +1,6 @@
-use crate::util::IntoOption;
-
+//! Private input types for the prover's secrets
 use super::{dlog_group, ProveDlog};
+use crate::util::IntoOption;
 use elliptic_curve::FromBytes;
 use k256::Scalar;
 
diff --git a/ergo-lib/src/sigma_protocol/prover.rs b/ergo-lib/src/sigma_protocol/prover.rs
index 462c1f4d3..1817447de 100644
--- a/ergo-lib/src/sigma_protocol/prover.rs
+++ b/ergo-lib/src/sigma_protocol/prover.rs
@@ -1,18 +1,22 @@
 //! Interpreter with enhanced functionality to prove statements.
 
+mod context_extension;
+mod prover_result;
+
+pub use context_extension::*;
+pub use prover_result::*;
+
 use super::{
     dlog_protocol,
     fiat_shamir::{fiat_shamir_hash_fn, fiat_shamir_tree_to_bytes},
+    private_input::PrivateInput,
     sig_serializer::serialize_sig,
     unchecked_tree::UncheckedSchnorr,
-    Challenge, PrivateInput, ProofTree, SigmaBoolean, SigmaProofOfKnowledgeTree,
-    UncheckedSigmaTree, UncheckedTree, UnprovenLeaf, UnprovenSchnorr, UnprovenTree,
-};
-use crate::{
-    chain::{context_extension::ContextExtension, prover_result::ProverResult},
-    eval::{Env, EvalError, Evaluator},
-    ErgoTree, ErgoTreeParsingError,
+    Challenge, ProofTree, SigmaBoolean, SigmaProofOfKnowledgeTree, UncheckedSigmaTree,
+    UncheckedTree, UnprovenLeaf, UnprovenSchnorr, UnprovenTree,
 };
+use crate::ergo_tree::{ErgoTree, ErgoTreeParsingError};
+use crate::eval::{Env, EvalError, Evaluator};
 use thiserror::Error;
 
 /// Prover errors
@@ -296,8 +300,7 @@ mod tests {
     use super::*;
     use crate::{
         ast::{Constant, ConstantVal, Expr},
-        chain::prover_result::ProofBytes,
-        sigma_protocol::DlogProverInput,
+        sigma_protocol::private_input::DlogProverInput,
         types::SType,
     };
     use std::rc::Rc;
diff --git a/ergo-lib/src/chain/context_extension.rs b/ergo-lib/src/sigma_protocol/prover/context_extension.rs
similarity index 100%
rename from ergo-lib/src/chain/context_extension.rs
rename to ergo-lib/src/sigma_protocol/prover/context_extension.rs
diff --git a/ergo-lib/src/chain/prover_result.rs b/ergo-lib/src/sigma_protocol/prover/prover_result.rs
similarity index 97%
rename from ergo-lib/src/chain/prover_result.rs
rename to ergo-lib/src/sigma_protocol/prover/prover_result.rs
index 93ae1d7f2..c1970def5 100644
--- a/ergo-lib/src/chain/prover_result.rs
+++ b/ergo-lib/src/sigma_protocol/prover/prover_result.rs
@@ -1,8 +1,7 @@
 //! ProverResult
 use std::io;
 
-use super::context_extension::ContextExtension;
-use super::{Base16DecodedBytes, Base16EncodedBytes};
+use crate::chain::{Base16DecodedBytes, Base16EncodedBytes};
 use crate::serialization::{
     sigma_byte_reader::SigmaByteRead, sigma_byte_writer::SigmaByteWrite, SerializationError,
     SigmaSerializable,
@@ -10,6 +9,8 @@ use crate::serialization::{
 #[cfg(feature = "json")]
 use serde::{Deserialize, Serialize};
 
+use super::ContextExtension;
+
 /// Serialized proof generated by ['Prover']
 #[cfg_attr(feature = "json", derive(Serialize, Deserialize))]
 #[cfg_attr(
diff --git a/ergo-lib/src/sigma_protocol/sig_serializer.rs b/ergo-lib/src/sigma_protocol/sig_serializer.rs
index d70e71b48..1f3f726d0 100644
--- a/ergo-lib/src/sigma_protocol/sig_serializer.rs
+++ b/ergo-lib/src/sigma_protocol/sig_serializer.rs
@@ -1,5 +1,6 @@
 //! Serialization of proof tree signatures
 
+use super::prover::ProofBytes;
 use super::{
     fiat_shamir::FiatShamirHash,
     unchecked_tree::{UncheckedLeaf, UncheckedSchnorr},
@@ -7,7 +8,6 @@ use super::{
     UncheckedTree,
 };
 
-use crate::chain::prover_result::ProofBytes;
 use k256::Scalar;
 use std::convert::{TryFrom, TryInto};
 
diff --git a/ergo-lib/src/sigma_protocol/verifier.rs b/ergo-lib/src/sigma_protocol/verifier.rs
index c06b85119..9bacf4829 100644
--- a/ergo-lib/src/sigma_protocol/verifier.rs
+++ b/ergo-lib/src/sigma_protocol/verifier.rs
@@ -1,5 +1,6 @@
 //! Verifier
 
+use super::prover::ProofBytes;
 use super::{
     dlog_protocol,
     fiat_shamir::{fiat_shamir_hash_fn, fiat_shamir_tree_to_bytes},
@@ -7,11 +8,8 @@ use super::{
     unchecked_tree::{UncheckedLeaf, UncheckedSchnorr},
     SigmaBoolean, UncheckedSigmaTree, UncheckedTree,
 };
-use crate::{
-    chain::prover_result::ProofBytes,
-    eval::{Env, EvalError, Evaluator},
-    ErgoTree, ErgoTreeParsingError,
-};
+use crate::ergo_tree::{ErgoTree, ErgoTreeParsingError};
+use crate::eval::{Env, EvalError, Evaluator};
 use dlog_protocol::FirstDlogProverMessage;
 
 /// Errors on proof verification
@@ -122,8 +120,8 @@ mod tests {
     use crate::{
         ast::{Constant, Expr},
         sigma_protocol::{
+            private_input::{DlogProverInput, PrivateInput},
             prover::{Prover, TestProver},
-            DlogProverInput, PrivateInput,
         },
         types::SType,
     };
diff --git a/ergo-lib/src/wallet.rs b/ergo-lib/src/wallet.rs
index 3eaa78a18..482f5566b 100644
--- a/ergo-lib/src/wallet.rs
+++ b/ergo-lib/src/wallet.rs
@@ -14,9 +14,10 @@ use crate::{
         ergo_box::ErgoBox, ergo_state_context::ErgoStateContext,
         transaction::unsigned::UnsignedTransaction, transaction::Transaction,
     },
-    sigma_protocol::prover::Prover,
-    sigma_protocol::prover::TestProver,
-    sigma_protocol::PrivateInput,
+    sigma_protocol::{
+        private_input::PrivateInput,
+        prover::{Prover, TestProver},
+    },
 };
 
 /// Wallet
diff --git a/ergo-lib/src/wallet/secret_key.rs b/ergo-lib/src/wallet/secret_key.rs
index d4659aa53..36fd51427 100644
--- a/ergo-lib/src/wallet/secret_key.rs
+++ b/ergo-lib/src/wallet/secret_key.rs
@@ -1,6 +1,6 @@
 //! Secret types
 use crate::chain::address::Address;
-use crate::sigma_protocol::{DlogProverInput, PrivateInput};
+use crate::sigma_protocol::private_input::{DlogProverInput, PrivateInput};
 
 /// Types of secrets
 #[derive(PartialEq, Debug, Clone)]
diff --git a/ergo-lib/src/wallet/signing.rs b/ergo-lib/src/wallet/signing.rs
index 3cae9e239..741ade0b7 100644
--- a/ergo-lib/src/wallet/signing.rs
+++ b/ergo-lib/src/wallet/signing.rs
@@ -1,10 +1,10 @@
 //! Transaction signing
 
+use crate::chain::transaction::Input;
 use crate::{
     chain::{
         ergo_box::ErgoBox,
         ergo_state_context::ErgoStateContext,
-        input::Input,
         transaction::{unsigned::UnsignedTransaction, Transaction},
     },
     eval::Env,
@@ -70,18 +70,19 @@ mod tests {
     use proptest::collection::vec;
     use proptest::prelude::*;
 
-    use crate::chain::ergo_box::box_builder::ErgoBoxCandidateBuilder;
-    use crate::chain::ergo_box::BoxValue;
     use crate::{
         ast::{Constant, Expr},
-        chain::{ergo_box::NonMandatoryRegisters, input::UnsignedInput, transaction::TxId},
+        chain::{
+            ergo_box::{box_builder::ErgoBoxCandidateBuilder, BoxValue, NonMandatoryRegisters},
+            transaction::{TxId, UnsignedInput},
+        },
+        ergo_tree::ErgoTree,
         sigma_protocol::{
+            private_input::{DlogProverInput, PrivateInput},
             prover::TestProver,
             verifier::{TestVerifier, Verifier, VerifierError},
-            DlogProverInput, PrivateInput,
         },
         types::SType,
-        ErgoTree,
     };
     use std::rc::Rc;
 
diff --git a/ergo-lib/src/wallet/tx_builder.rs b/ergo-lib/src/wallet/tx_builder.rs
index 93e28094b..85bd8e5a0 100644
--- a/ergo-lib/src/wallet/tx_builder.rs
+++ b/ergo-lib/src/wallet/tx_builder.rs
@@ -4,32 +4,20 @@ use std::collections::HashSet;
 
 use thiserror::Error;
 
-use crate::chain::address::Address;
-use crate::chain::address::AddressEncoder;
-use crate::chain::address::NetworkPrefix;
+use crate::chain::address::{Address, AddressEncoder, NetworkPrefix};
 use crate::chain::contract::Contract;
-use crate::chain::data_input::DataInput;
-use crate::chain::ergo_box::box_builder::ErgoBoxCandidateBuilder;
-use crate::chain::ergo_box::box_builder::ErgoBoxCandidateBuilderError;
-use crate::chain::ergo_box::sum_tokens_from_boxes;
-use crate::chain::ergo_box::sum_value;
-use crate::chain::ergo_box::BoxId;
-use crate::chain::ergo_box::BoxValue;
-use crate::chain::ergo_box::BoxValueError;
-use crate::chain::input::Input;
-use crate::chain::prover_result::ProofBytes;
-use crate::chain::prover_result::ProverResult;
-use crate::chain::token::Token;
-use crate::chain::token::TokenId;
-use crate::chain::transaction::Transaction;
+use crate::chain::ergo_box::box_builder::{ErgoBoxCandidateBuilder, ErgoBoxCandidateBuilderError};
+use crate::chain::ergo_box::{sum_tokens_from_boxes, sum_value, BoxId, BoxValue, BoxValueError};
+use crate::chain::token::{Token, TokenId};
+use crate::chain::transaction::{DataInput, Input, Transaction, UnsignedInput};
 use crate::chain::{
-    ergo_box::ErgoBoxAssets, ergo_box::ErgoBoxCandidate, ergo_box::ErgoBoxId, input::UnsignedInput,
+    ergo_box::ErgoBoxAssets, ergo_box::ErgoBoxCandidate, ergo_box::ErgoBoxId,
     transaction::unsigned::UnsignedTransaction,
 };
 use crate::constants::MINERS_FEE_MAINNET_ADDRESS;
-use crate::serialization::SerializationError;
-use crate::serialization::SigmaSerializable;
+use crate::serialization::{SerializationError, SigmaSerializable};
 use crate::sigma_protocol;
+use crate::sigma_protocol::prover::{ProofBytes, ProverResult};
 
 use super::box_selector::{BoxSelection, BoxSelectorError};
 
@@ -254,19 +242,14 @@ mod tests {
 
     use proptest::{collection::vec, prelude::*};
 
-    use crate::chain::ergo_box::checked_sum;
-    use crate::chain::ergo_box::ErgoBox;
-    use crate::chain::ergo_box::NonMandatoryRegisters;
-    use crate::chain::token::tests::ArbTokenIdParam;
-    use crate::chain::token::Token;
-    use crate::chain::token::TokenAmount;
-    use crate::chain::token::TokenId;
-    use crate::chain::transaction::TxId;
-    use crate::test_util::force_any_val;
-    use crate::test_util::force_any_val_with;
-    use crate::wallet::box_selector::BoxSelector;
-    use crate::wallet::box_selector::SimpleBoxSelector;
-    use crate::ErgoTree;
+    use crate::chain::{
+        ergo_box::{checked_sum, ErgoBox, NonMandatoryRegisters},
+        token::{tests::ArbTokenIdParam, Token, TokenAmount, TokenId},
+        transaction::TxId,
+    };
+    use crate::ergo_tree::ErgoTree;
+    use crate::test_util::{force_any_val, force_any_val_with};
+    use crate::wallet::box_selector::{BoxSelector, SimpleBoxSelector};
 
     use super::*;
 


Test patch: