@raydium-io/raydium-sdk-v2/src/marshmallow/index.ts

An SDK for building applications on top of Raydium.

import { PublicKey } from "@solana/web3.js";
import BN, { isBN } from "bn.js";

import {
  bits,
  blob,
  Blob,
  Layout,
  offset as _offset,
  seq as _seq,
  Structure as _Structure,
  u32 as _u32,
  u8 as _u8,
  UInt,
  union as _union,
  Union as _Union,
} from "./buffer-layout";

export * from "./buffer-layout";
export { blob };

export class BNLayout<P extends string = ""> extends Layout<BN, P> {
  blob: Layout<Buffer>;
  signed: boolean;

  constructor(span: number, signed: boolean, property?: P) {
    //@ts-expect-error type wrong for super()'s type different from extends, but it desn't matter
    super(span, property);
    this.blob = blob(span);
    this.signed = signed;
  }

  /** @override */
  decode(b: Buffer, offset = 0): BN {
    const num = new BN(this.blob.decode(b, offset), 10, "le");
    if (this.signed) {
      return num.fromTwos(this.span * 8).clone();
    }
    return num;
  }

  /** @override */
  encode(src: BN, b: Buffer, offset = 0): number {
    if (typeof src === "number") src = new BN(src); // src will pass a number accidently in union
    if (this.signed) {
      src = src.toTwos(this.span * 8);
    }
    return this.blob.encode(src.toArrayLike(Buffer, "le", this.span), b, offset);
  }
}

export class WideBits<P extends string = ""> extends Layout<Record<string, boolean>, P> {
  _lower: any;
  _upper: any;
  // TODO: unknown
  constructor(property?: P) {
    //@ts-expect-error type wrong for super()'s type different from extends , but it desn't matter
    super(8, property);
    this._lower = bits(_u32(), false);
    this._upper = bits(_u32(), false);
  }

  addBoolean(property: string): void {
    if (this._lower.fields.length < 32) {
      this._lower.addBoolean(property);
    } else {
      this._upper.addBoolean(property);
    }
  }

  decode(b: Buffer, offset = 0): Record<string, boolean> {
    const lowerDecoded = this._lower.decode(b, offset);
    const upperDecoded = this._upper.decode(b, offset + this._lower.span);
    return { ...lowerDecoded, ...upperDecoded };
  }

  encode(src: any /* TEMP */, b: Buffer, offset = 0): any {
    return this._lower.encode(src, b, offset) + this._upper.encode(src, b, offset + this._lower.span);
  }
}

export function u8<P extends string = "">(property?: P): UInt<number, P> {
  return new UInt(1, property);
}

export function u32<P extends string = "">(property?: P): UInt<number, P> {
  return new UInt(4, property);
}

export function u64<P extends string = "">(property?: P): BNLayout<P> {
  return new BNLayout(8, false, property);
}

export function u128<P extends string = "">(property?: P): BNLayout<P> {
  return new BNLayout(16, false, property);
}

export function i8<P extends string = "">(property?: P): BNLayout<P> {
  return new BNLayout(1, true, property);
}

export function i64<P extends string = "">(property?: P): BNLayout<P> {
  return new BNLayout(8, true, property);
}

export function i128<P extends string = "">(property?: P): BNLayout<P> {
  return new BNLayout(16, true, property);
}

export class WrappedLayout<T, U, P extends string = ""> extends Layout<U, P> {
  layout: Layout<T>;
  decoder: (data: T) => U;
  encoder: (src: U) => T;

  constructor(layout: Layout<T>, decoder: (data: T) => U, encoder: (src: U) => T, property?: P) {
    //@ts-expect-error type wrong for super()'s type different from extends , but it desn't matter
    super(layout.span, property);
    this.layout = layout;
    this.decoder = decoder;
    this.encoder = encoder;
  }

  decode(b: Buffer, offset?: number): U {
    return this.decoder(this.layout.decode(b, offset));
  }

  encode(src: U, b: Buffer, offset?: number): number {
    return this.layout.encode(this.encoder(src), b, offset);
  }

  getSpan(b: Buffer, offset?: number): number {
    return this.layout.getSpan(b, offset);
  }
}

export function publicKey<P extends string = "">(property?: P): Layout<PublicKey, P> {
  return new WrappedLayout(
    blob(32),
    (b: Buffer) => new PublicKey(b),
    (key: PublicKey) => key.toBuffer(),
    property,
  );
}

export class OptionLayout<T, P> extends Layout<T | null, P> {
  layout: Layout<T>;
  discriminator: Layout<number>;

  constructor(layout: Layout<T>, property?: P) {
    //@ts-expect-error type wrong for super()'s type different from extends , but it desn't matter
    super(-1, property);
    this.layout = layout;
    this.discriminator = _u8();
  }

  encode(src: T | null, b: Buffer, offset = 0): number {
    if (src === null || src === undefined) {
      return this.discriminator.encode(0, b, offset);
    }
    this.discriminator.encode(1, b, offset);
    return this.layout.encode(src, b, offset + 1) + 1;
  }

  decode(b: Buffer, offset = 0): T | null {
    const discriminator = this.discriminator.decode(b, offset);
    if (discriminator === 0) {
      return null;
    } else if (discriminator === 1) {
      return this.layout.decode(b, offset + 1);
    }
    throw new Error("Invalid option " + this.property);
  }

  getSpan(b: Buffer, offset = 0): number {
    const discriminator = this.discriminator.decode(b, offset);
    if (discriminator === 0) {
      return 1;
    } else if (discriminator === 1) {
      return this.layout.getSpan(b, offset + 1) + 1;
    }
    throw new Error("Invalid option " + this.property);
  }
}

export function option<T, P extends string = "">(layout: Layout<T>, property?: P): Layout<T | null, P> {
  return new OptionLayout<T, P>(layout, property);
}

export function bool<P extends string = "">(property?: P): Layout<boolean, P> {
  return new WrappedLayout(_u8(), decodeBool, encodeBool, property);
}

export function decodeBool(value: number): boolean {
  if (value === 0) {
    return false;
  } else if (value === 1) {
    return true;
  }
  throw new Error("Invalid bool: " + value);
}

export function encodeBool(value: boolean): number {
  return value ? 1 : 0;
}

export function vec<T, P extends string = "">(elementLayout: Layout<T>, property?: P): Layout<T[], P> {
  const length = _u32("length");
  const layout: Layout<{ values: T[] }> = struct([
    length,
    seq(elementLayout, _offset(length, -length.span), "values"),
  ]) as any; // Something I don't know
  return new WrappedLayout(
    layout,
    ({ values }) => values,
    (values) => ({ values }),
    property,
  );
}

export function tagged<T, P extends string = "">(tag: BN, layout: Layout<T>, property?: P): Layout<T, P> {
  const wrappedLayout: Layout<{ tag: BN; data: T }> = struct([u64("tag"), layout.replicate("data")]) as any; // Something I don't know

  function decodeTag({ tag: receivedTag, data }: { tag: BN; data: T }): T {
    if (!receivedTag.eq(tag)) {
      throw new Error("Invalid tag, expected: " + tag.toString("hex") + ", got: " + receivedTag.toString("hex"));
    }
    return data;
  }

  return new WrappedLayout(wrappedLayout, decodeTag, (data) => ({ tag, data }), property);
}

export function vecU8<P extends string = "">(property?: P): Layout<Buffer, P> {
  const length = _u32("length");
  const layout: Layout<{ data: Buffer }> = struct([length, blob(_offset(length, -length.span), "data")]) as any; // Something I don't know
  return new WrappedLayout(
    layout,
    ({ data }) => data,
    (data) => ({ data }),
    property,
  );
}

export function str<P extends string = "">(property?: P): Layout<string, P> {
  return new WrappedLayout(
    vecU8(),
    (data) => data.toString("utf-8"),
    (s) => Buffer.from(s, "utf-8"),
    property,
  );
}

export interface EnumLayout<T, P extends string = ""> extends Layout<T, P> {
  registry: Record<string, Layout<any>>;
}

export function rustEnum<T, P extends string = "">(variants: Layout<any>[], property?: P): EnumLayout<T, P> {
  const unionLayout = _union(_u8(), property);
  variants.forEach((variant, index) => unionLayout.addVariant(index, variant, variant.property));
  return unionLayout as any; // ?why use UnionLayout? This must be a fault
}

export function array<T, P extends string = "">(
  elementLayout: Layout<T>,
  length: number,
  property?: P,
): Layout<T[], P> {
  const layout = struct([seq(elementLayout, length, "values")]) as any as Layout<{ values: T[] }>; // Something I don't know
  return new WrappedLayout(
    layout,
    ({ values }) => values,
    (values) => ({ values }),
    property,
  );
}

export class Structure<T, P, D extends { [key: string]: any; }> extends _Structure<T, P, D> {
  /** @override */
  decode(b: Buffer, offset?: number): D {
    return super.decode(b, offset);
  }
}

export function struct<T, P extends string = "">(
  fields: T,
  property?: P,
  decodePrefixes?: boolean,
): T extends Layout<infer Value, infer Property>[]
  ? Structure<
    Value,
    P,
    {
      [K in Exclude<Extract<Property, string>, "">]: Extract<T[number], Layout<any, K>> extends Layout<infer V, any>
      ? V
      : any;
    }
  >
  : any {
  //@ts-expect-error this type is not quite satisfied the define, but, never no need to worry about.
  return new Structure(fields, property, decodePrefixes);
}

export type GetLayoutSchemaFromStructure<T extends Structure<any, any, any>> = T extends Structure<any, any, infer S>
  ? S
  : any;
export type GetStructureFromLayoutSchema<S extends { [key: string]: any; }> = Structure<any, any, S>;

export class Union<Schema extends { [key: string]: any; }> extends _Union<Schema> {
  encodeInstruction(instruction: any): Buffer {
    const instructionMaxSpan = Math.max(...Object.values(this.registry).map((r) => r.span));
    const b = Buffer.alloc(instructionMaxSpan);
    return b.slice(0, this.encode(instruction, b));
  }

  decodeInstruction(instruction: any): Partial<Schema> {
    return this.decode(instruction);
  }
}
export function union<UnionSchema extends { [key: string]: any } = any>(
  discr: any,
  defaultLayout?: any,
  property?: string,
): Union<UnionSchema> {
  return new Union(discr, defaultLayout, property);
}

class Zeros extends Blob {
  decode(b: Buffer, offset: number): Buffer {
    const slice = super.decode(b, offset);
    if (!slice.every((v) => v === 0)) {
      throw new Error("nonzero padding bytes");
    }
    return slice;
  }
}

export function zeros(length: number): Zeros {
  return new Zeros(length);
}

export function seq<T, P extends string = "", AnotherP extends string = "">(
  elementLayout: Layout<T, P>,
  count: number | BN | Layout<BN | number, P>,
  property?: AnotherP,
): Layout<T[], AnotherP> {
  let parsedCount: number;
  const superCount =
    typeof count === "number"
      ? count
      : isBN(count)
        ? count.toNumber()
        : new Proxy(count as unknown as Layout<number> /* pretend to be Layout<number> */, {
          get(target, property): any {
            if (!parsedCount) {
              // get count in targetLayout. note that count may be BN
              const countProperty = Reflect.get(target, "count");

              // let targetLayout's  property:count be a number
              parsedCount = isBN(countProperty) ? countProperty.toNumber() : countProperty;

              // record the count
              Reflect.set(target, "count", parsedCount);
            }
            return Reflect.get(target, property);
          },
          set(target, property, value): any {
            if (property === "count") {
              parsedCount = value;
            }
            return Reflect.set(target, property, value);
          },
        });

  // @ts-expect-error force type
  return _seq(elementLayout, superCount, property);
}