Add ErgoTree constant replacing and template extraction

.changeset/cruel-parks-prove.md

@@ -0,0 +1,5 @@
+---
+"@fleet-sdk/serializer": patch
+---
+
+Memoize `SConstant` bytes for efficient re-encoding

.changeset/silver-rabbits-shake.md

@@ -0,0 +1,5 @@
+---
+"@fleet-sdk/serializer": patch
+---
+
+Add `SigmaByteReader#readRemainingBytes` method

.changeset/tough-jars-sleep.md

@@ -0,0 +1,5 @@
+---
+"@fleet-sdk/compiler": patch
+---
+
+Replace `CompilerOutput` with `ErgoTree`

packages/compiler/src/compiler.spec.ts

@@ -17,6 +17,13 @@ const compilerTestVectors: {
   template: string;
   options: CompilerOptions;
 }[] = [
+  {
+    name: "v0 - Segregated constants",
+    script: "sigmaProp(HEIGHT > 100 && HEIGHT < 200)",
+    tree: "100204c801049003d1ed91a373008fa37301",
+    template: "d1ed91a373008fa37301",
+    options: { version: 0, segregateConstants: true, includeSize: false }
+  },
   {
     name: "v0 - Segregated constants",
     script: "sigmaProp(HEIGHT > 100)",
@@ -89,8 +96,8 @@ describe("ErgoScript Compiler", () => {
     const tree = compile(tv.script, tv.options);
 
     expect(tree.toHex()).to.be.equal(tv.tree);
-    expect(tree.template.toHex()).to.be.equal(tv.template);
-    expect(hex.encode(tree.template.toBytes())).to.be.equal(tv.template);
+    expect(hex.encode(tree.template)).to.be.equal(tv.template);
+    expect(hex.encode(tree.template)).to.be.equal(tv.template);
 
     expect(tree.hasSegregatedConstants).to.be.equal(tv.options.segregateConstants);
     expect(tree.version).to.be.equal(tv.options.version);

packages/compiler/src/compiler.ts

@@ -6,14 +6,14 @@ import {
   isEmpty,
   isHex
 } from "@fleet-sdk/common";
+import { ErgoTree } from "@fleet-sdk/core";
 import { SConstant } from "@fleet-sdk/serializer";
 import {
   SigmaCompiler$,
   type SigmaCompilerNamedConstantsMap,
   Value,
   Value$
 } from "sigmastate-js/main";
-import { CompilerOutput } from "./compilerOutput";
 
 export type NamedConstantsMap = {
   [key: string]: string | Value | SConstant;
@@ -74,7 +74,7 @@ export const compilerDefaults: Required<CompilerOptions> = {
  * @param options - Optional compiler options to customize the compilation process.
  * @returns The output of the compilation process.
  */
-export function compile(script: string, options?: CompilerOptions): CompilerOutput {
+export function compile(script: string, options?: CompilerOptions): ErgoTree {
   const opt = ensureDefaults(options, compilerDefaults);
   assert(opt.version < 8, `Version should be lower than 8, got ${opt.version}`);
 
@@ -86,14 +86,13 @@ export function compile(script: string, options?: CompilerOptions): CompilerOutp
   const compiler =
     opt.network === "mainnet" ? SigmaCompiler$.forMainnet() : SigmaCompiler$.forTestnet();
 
-  const tree = compiler.compile(
-    parseNamedConstantsMap(opt.map),
-    opt.segregateConstants,
-    headerFlags,
-    script
+  const output = Uint8Array.from(
+    compiler
+      .compile(parseNamedConstantsMap(opt.map), opt.segregateConstants, headerFlags, script)
+      .bytes().u
   );
 
-  return new CompilerOutput(tree, opt.network === "mainnet" ? Network.Mainnet : Network.Testnet);
+  return new ErgoTree(output, opt.network === "mainnet" ? Network.Mainnet : Network.Testnet);
 }
 
 export function parseNamedConstantsMap(map: NamedConstantsMap): SigmaCompilerNamedConstantsMap {

packages/compiler/src/index.ts

@@ -1,3 +1 @@
 export { compile } from "./compiler";
-export * from "./compilerOutput";
-export * from "./contractTemplate";

packages/core/src/models/ergoTree.spec.ts

@@ -1,6 +1,9 @@
 import { Network } from "@fleet-sdk/common";
 import { hex } from "@fleet-sdk/crypto";
+import { SBool, SigmaByteReader, estimateVLQSize } from "@fleet-sdk/serializer";
+import { ErgoTree$ } from "sigmastate-js/main";
 import { describe, expect, it, test } from "vitest";
+import { SInt } from "../constantSerializer";
 import { ErgoAddress } from "./ergoAddress";
 import { ErgoTree } from "./ergoTree";
 
@@ -59,6 +62,185 @@ describe("ErgoTree model", () => {
   });
 });
 
+describe("Constant handling", () => {
+  it("Should parse constants", () => {
+    const treeHex = "18090104c801d191a37300";
+
+    const tree = new ErgoTree(treeHex);
+    expect(tree.constants).not.to.be.empty;
+    expect(tree.constants).to.have.length(1);
+    expect(tree.constants?.[0].data).to.be.equal(100);
+
+    const sigmaJsTree = ErgoTree$.fromHex(treeHex);
+    expect(sigmaJsTree.constants()).not.to.be.undefined;
+    expect(sigmaJsTree.constants()).to.have.length(1);
+    expect(sigmaJsTree.constants()[0].data).to.be.equal(100);
+
+    // todo: move to template generation specific test
+    expect(sigmaJsTree.templateHex()).to.be.equal(hex.encode(tree.template));
+  });
+
+  it("should replace constant", () => {
+    const originalTreeHex = "100104c801d191a37300"; // sigmaProp(HEIGHT > 100)
+    const originalTree = new ErgoTree(originalTreeHex);
+
+    expect(originalTree.constants).to.have.length(1);
+    expect(originalTree.constants?.[0].data).to.be.equal(100);
+
+    originalTree.replaceConstant(0, SInt(200));
+    expect(originalTree.constants?.[0].data).to.be.equal(200);
+
+    const patchedTreeHex = originalTree.toHex();
+    const patchedTree = new ErgoTree(patchedTreeHex);
+    expect(originalTree.template).to.be.deep.equal(patchedTree.template);
+    expect(patchedTree.constants).to.have.length(1);
+    expect(patchedTree.constants?.[0].data).to.be.equal(200);
+  });
+
+  it("Should replace constant with a bigger one", () => {
+    const originalTreeHex = "18090104c801d191a37300"; // sigmaProp(HEIGHT > 100), 100 needs 1 byte
+    const originalTree = new ErgoTree(originalTreeHex);
+    expect(originalTree.hasSize).to.be.true;
+    expect(originalTree.constants).to.have.length(1);
+    expect(originalTree.constants?.[0].data).to.be.equal(100);
+
+    let reader = new SigmaByteReader(originalTreeHex);
+    reader.readByte(); // skip header
+    let size = reader.readUInt();
+    const originalTreeSize = 1 + estimateVLQSize(size) + size; // header + vlq size + body
+    expect(originalTree.bytes.length).to.be.equal(originalTreeSize); // header + vlq size + body
+
+    originalTree.replaceConstant(0, SInt(20000)); // 20000 needs 3 bytes
+    expect(originalTree.constants?.[0].data).to.be.equal(20000);
+
+    const patchedTreeHex = originalTree.toHex();
+    const patchedTree = new ErgoTree(patchedTreeHex);
+    expect(originalTree.template).to.be.deep.equal(patchedTree.template);
+    expect(patchedTree.constants).to.have.length(1);
+    expect(patchedTree.constants?.[0].data).to.be.equal(20000);
+
+    // check if the size is updated
+    reader = new SigmaByteReader(patchedTreeHex);
+    reader.readByte(); // skip header
+    size = reader.readUInt();
+    const patchedTreeSize = 1 + estimateVLQSize(size) + size; // header + vlq size + body
+    expect(patchedTree.bytes.length).to.be.equal(patchedTreeSize); // header + vlq size + body
+    expect(patchedTreeSize).to.be.greaterThan(
+      originalTreeSize,
+      "original tree size should be smaller"
+    );
+  });
+
+  it("Should replace constant and preserve others", () => {
+    const tree = new ErgoTree("100204c801049003d1ed91a373008fa37301"); // sigmaProp(HEIGHT > 100 && HEIGHT < 200)
+    expect(tree.hasSegregatedConstants).to.be.true;
+    expect(tree.constants).to.have.length(2);
+    expect(tree.constants?.[0].data).to.be.equal(100);
+    expect(tree.constants?.[1].data).to.be.equal(200);
+
+    tree.replaceConstant(0, SInt(300));
+
+    expect(tree.constants).to.have.length(2);
+    expect(tree.constants?.[0].data).to.be.equal(300);
+    expect(tree.constants?.[1].data).to.be.equal(200);
+
+    const reconstructedTree = new ErgoTree(tree.bytes);
+    expect(reconstructedTree.constants).to.have.length(2);
+    expect(reconstructedTree.constants?.[0].data).to.be.equal(300);
+    expect(reconstructedTree.constants?.[0].type.toString()).to.be.equal("SInt");
+    expect(reconstructedTree.constants?.[1].data).to.be.equal(200);
+    expect(reconstructedTree.constants?.[1].type.toString()).to.be.equal("SInt");
+  });
+
+  it("should fail to replace constant in a non-segregated tree", () => {
+    const tree = new ErgoTree("00d191a304c801");
+    expect(tree.hasSegregatedConstants).to.be.false;
+    expect(() => tree.replaceConstant(0, SInt(200))).to.throw(
+      "Constant segregation is not enabled."
+    );
+  });
+
+  it("should fail to replace constant of different type", () => {
+    const tree = new ErgoTree("18090104c801d191a37300"); // sigmaProp(HEIGHT > 100)
+    expect(tree.hasSegregatedConstants).to.be.true;
+    expect(tree.constants).to.have.length(1);
+    expect(() => tree.replaceConstant(0, SBool(true))).to.throw(
+      "Constant type mismatch: can't replace 'SInt' with 'SBool'"
+    );
+  });
+
+  it("should fail to replace non-existing constant", () => {
+    const tree = new ErgoTree("18090104c801d191a37300"); // sigmaProp(HEIGHT > 100)
+    expect(tree.hasSegregatedConstants).to.be.true;
+    expect(tree.constants).to.have.length(1);
+    expect(() => tree.replaceConstant(1, SInt(200))).to.throw("Constant at index 1 not found.");
+  });
+});
+
+describe("Serialization", () => {
+  it("Should parse and reconstruct tree with segregated constants", () => {
+    const treeHex = "18090104c801d191a37300"; // sigmaProp(HEIGHT > 100)
+    const tree = new ErgoTree(treeHex);
+    expect(tree.hasSegregatedConstants).to.be.true;
+    expect(tree.constants).not.to.be.empty; // checking constants triggers parsing
+    expect(tree.constants).to.have.length(1);
+    expect(tree.constants?.[0].data).to.be.equal(100);
+    expect(tree.toHex()).to.be.equal(treeHex); // check if the tree is reconstructed correctly
+  });
+
+  it("should parse and reconstruct tree with non-segregated consts", () => {
+    const treeHex = "00d191a304c801"; // sigmaProp(HEIGHT > 100)
+    const tree = new ErgoTree(treeHex);
+    expect(tree.hasSegregatedConstants).to.be.false;
+    expect(tree.constants).to.be.empty; // checking constants don't trigger parsing if `hasSegregatedConstants == false`
+    expect(tree.template).not.to.be.empty; // checking template triggers parsing
+    expect(tree.toHex()).to.be.equal(treeHex); // check if the tree is reconstructed correctly
+  });
+
+  it("Should parse and reconstruct tree with segregated constants, but no constants", () => {
+    const treeHex = "1000d17300"; // sigmaProp(false), manually tweaked to have empty constants
+    const tree = new ErgoTree(treeHex);
+    expect(tree.hasSegregatedConstants).to.be.true;
+    expect(tree.constants).to.be.empty;
+    expect(tree.toHex()).to.be.equal(treeHex); // check if the tree is reconstructed correctly
+  });
+
+  it("should fail to parse empty root tree", () => {
+    const tree = new ErgoTree("00");
+    expect(tree.hasSegregatedConstants).to.be.false;
+    expect(() => tree.template /* triggers parsing */).to.throw("Empty tree bytes.");
+  });
+
+  it("Should ignore pushing new constants to a tree", () => {
+    const treeHex = "18090104c801d191a37300"; // sigmaProp(HEIGHT > 100)
+    const tree = new ErgoTree(treeHex);
+    expect(tree.hasSegregatedConstants).to.be.true;
+    expect(tree.constants).to.have.length(1);
+
+    // @ts-expect-error constants prop is readonly
+    tree.constants.push(SInt(200));
+
+    expect(tree.constants).to.have.length(1); // should not change the length
+    expect(tree.toHex()).to.be.equal(treeHex); // check if the tree is reconstructed correctly
+  });
+});
+
+describe("Template generation", () => {
+  it("should ignore direct template changes", () => {
+    const treeHex = "18090104c801d191a37300"; // sigmaProp(HEIGHT > 100)
+    const tree = new ErgoTree(treeHex);
+
+    // @ts-expect-error template prop is readonly
+    tree.template[0] = 0xff; // change first byte to 0xff
+
+    expect(tree.template[0]).not.to.be.equal(0xff); // check if the template is changed
+
+    tree.replaceConstant(0, SInt(100)); // trigger re-serializations and should not change the template
+
+    expect(tree.toHex()).to.be.equal(treeHex); // check if the tree is reconstructed correctly
+  });
+});
+
 describe("Encoding", () => {
   const treeHex = "100104c801d191a37300";