@gmod/gbz-base 0.0.1 → 1.1.0

package/src/db.ts

@@ -0,0 +1,374 @@
+import { ENDMARKER, nodeId, nodeOrientation } from './gbwt/node.ts'
+import { GbwtRecord, decompressEdges } from './gbwt/record.ts'
+import { decodeSequence, encodedSequenceLength } from './gbwt/sequence.ts'
+import { graphNameFromTags } from './graphName.ts'
+import { SqliteDatabase } from './sqlite/database.ts'
+
+import type { ByteSource } from './filehandle.ts'
+import type { Pos } from './gbwt/record.ts'
+import type { GraphName } from './graphName.ts'
+import type { PagerOptions } from './sqlite/pager.ts'
+import type { SqlValue } from './sqlite/record.ts'
+
+export interface PathName {
+  sample: string
+  contig: string
+  haplotype: number
+  fragment: number
+}
+
+export const GENERIC_SAMPLE = '_gbwt_ref'
+export const SCHEMA_VERSION = 'GBZ-base version 4'
+
+export class SchemaVersionError extends Error {
+  override name = 'SchemaVersionError'
+
+  readonly found: string | undefined
+
+  constructor(found: string | undefined) {
+    super(
+      found === undefined
+        ? `not a gbz-base database: its Tags table has no version`
+        : `unsupported database schema "${found}"; this reader understands "${SCHEMA_VERSION}"`,
+    )
+    this.found = found
+  }
+}
+
+export function formatPathName(name: PathName, end: number) {
+  return `${name.sample}#${name.haplotype}#${name.contig}[${name.fragment}-${end}]`
+}
+
+export interface HaplotypeSample {
+  node: number
+  offset: number
+  pathHandle: number
+  orientation: 'forward' | 'reverse'
+  pathOffset: number
+}
+
+export interface GbzPath {
+  handle: number
+  fwStart: Pos
+  revStart: Pos
+  name: PathName
+  isIndexed: boolean
+}
+
+export class GbzRecord {
+  readonly sequenceLen: number
+  readonly handle: number
+  readonly edges: Pos[]
+  readonly bwt: Uint8Array
+  readonly encodedSequence: Uint8Array
+  readonly next: number | undefined
+  private decoded: string | undefined
+
+  constructor(
+    handle: number,
+    edges: Pos[],
+    bwt: Uint8Array,
+    encodedSequence: Uint8Array,
+    next: number | undefined,
+  ) {
+    this.handle = handle
+    this.edges = edges
+    this.bwt = bwt
+    this.encodedSequence = encodedSequence
+    this.next = next
+    this.sequenceLen = encodedSequenceLength(encodedSequence)
+  }
+
+  get id() {
+    return nodeId(this.handle)
+  }
+
+  get orientation() {
+    return nodeOrientation(this.handle)
+  }
+
+  get sequence() {
+    this.decoded ??= decodeSequence(this.encodedSequence)
+    return this.decoded
+  }
+
+  successors() {
+    return this.edges.filter(e => e.node !== ENDMARKER).map(e => e.node)
+  }
+
+  gbwt() {
+    if (this.edges.length === 0) {
+      throw new Error(`GBWT record for handle ${this.handle} is empty`)
+    }
+    return new GbwtRecord(this.edges, this.bwt)
+  }
+}
+
+function num(value: SqlValue | undefined, what: string) {
+  if (typeof value !== 'number') {
+    throw new Error(`${what} is not a number in the database`)
+  }
+  return value
+}
+
+function str(value: SqlValue | undefined, what: string) {
+  if (typeof value !== 'string') {
+    throw new Error(`${what} is not text in the database`)
+  }
+  return value
+}
+
+function blob(value: SqlValue | undefined, what: string) {
+  if (!(value instanceof Uint8Array)) {
+    throw new Error(`${what} is not a blob in the database`)
+  }
+  return value
+}
+
+function rowToPath(rowid: number, values: SqlValue[]): GbzPath {
+  return {
+    handle: rowid,
+    fwStart: {
+      node: num(values[1], 'Paths.fw_node'),
+      offset: num(values[2], 'Paths.fw_offset'),
+    },
+    revStart: {
+      node: num(values[3], 'Paths.rev_node'),
+      offset: num(values[4], 'Paths.rev_offset'),
+    },
+    name: {
+      sample: str(values[5], 'Paths.sample'),
+      contig: str(values[6], 'Paths.contig'),
+      haplotype: num(values[7], 'Paths.haplotype'),
+      fragment: num(values[8], 'Paths.fragment'),
+    },
+    isIndexed: num(values[9], 'Paths.is_indexed') !== 0,
+  }
+}
+
+export interface OpenOptions extends PagerOptions {
+  haplotypeIndex?: ByteSource
+}
+
+async function readTags(sqlite: SqliteDatabase) {
+  const tags = new Map<string, string>()
+  for await (const { values } of sqlite.scan('Tags')) {
+    tags.set(str(values[0], 'Tags.key'), str(values[1], 'Tags.value'))
+  }
+  return tags
+}
+
+export class GBZBase {
+  private tagCache: Promise<Map<string, string>> | undefined
+  private pathCache: Promise<GbzPath[]> | undefined
+  private indexTags: Map<string, string> | undefined
+
+  readonly sqlite: SqliteDatabase
+  private readonly index: SqliteDatabase
+
+  private constructor(sqlite: SqliteDatabase, index: SqliteDatabase) {
+    this.sqlite = sqlite
+    this.index = index
+  }
+
+  static async open(source: ByteSource, opts: OpenOptions = {}) {
+    const { haplotypeIndex, ...pagerOptions } = opts
+    const sqlite = await SqliteDatabase.open(source, pagerOptions)
+    for (const table of ['Tags', 'Nodes', 'Paths', 'ReferenceIndex']) {
+      sqlite.rootPage(table)
+    }
+    const index = haplotypeIndex
+      ? await SqliteDatabase.open(haplotypeIndex, pagerOptions)
+      : sqlite
+    const db = new GBZBase(sqlite, index)
+    const version = await db.tag('version')
+    if (version !== SCHEMA_VERSION) {
+      throw new SchemaVersionError(version)
+    }
+    if (haplotypeIndex) {
+      for (const table of ['Tags', 'HaplotypeSamples', 'HaplotypeLengths']) {
+        index.rootPage(table)
+      }
+      db.indexTags = await readTags(index)
+      const indexed = db.indexTags.get('haplotype_index_paths')
+      const paths = await db.tag('paths')
+      if (indexed !== paths) {
+        throw new Error(
+          `haplotype index was built for ${indexed ?? 'an unknown number of'} paths but the graph has ${paths}`,
+        )
+      }
+    }
+    return db
+  }
+
+  tags() {
+    this.tagCache ??= readTags(this.sqlite)
+    return this.tagCache
+  }
+
+  async tag(key: string) {
+    return (await this.tags()).get(key)
+  }
+
+  async getRecord(handle: number) {
+    const row = await this.sqlite.byRowid('Nodes', handle)
+    if (!row) {
+      return undefined
+    }
+    const next = row[4]
+    return new GbzRecord(
+      handle,
+      decompressEdges(blob(row[1], 'Nodes.edges')),
+      blob(row[2], 'Nodes.bwt'),
+      blob(row[3], 'Nodes.sequence'),
+      typeof next === 'number' ? next : undefined,
+    )
+  }
+
+  paths() {
+    this.pathCache ??= (async () => {
+      const paths: GbzPath[] = []
+      for await (const { rowid, values } of this.sqlite.scan('Paths')) {
+        paths.push(rowToPath(rowid, values))
+      }
+      return paths
+    })()
+    return this.pathCache
+  }
+
+  async getPath(handle: number) {
+    const row = await this.sqlite.byRowid('Paths', handle)
+    return row ? rowToPath(handle, row) : undefined
+  }
+
+  async findPath(name: PathName) {
+    const candidates = (await this.paths()).filter(
+      p =>
+        p.name.sample === name.sample &&
+        p.name.contig === name.contig &&
+        p.name.haplotype === name.haplotype &&
+        p.name.fragment <= name.fragment,
+    )
+    return candidates.sort((a, b) => b.name.fragment - a.name.fragment)[0]
+  }
+
+  async pathsForSample(sample: string) {
+    return (await this.paths()).filter(p => p.name.sample === sample)
+  }
+
+  async graphName() {
+    const gbzTags = new Map<string, string>()
+    for (const [key, value] of await this.tags()) {
+      if (key.startsWith('gbz_')) {
+        gbzTags.set(key.slice('gbz_'.length), value)
+      }
+    }
+    let name: GraphName = {
+      name: undefined,
+      subgraph: new Map(),
+      translation: new Map(),
+    }
+    try {
+      name = graphNameFromTags(gbzTags)
+    } catch {
+      // upstream falls back to an empty name when the tags do not parse
+    }
+    return name
+  }
+
+  async hasChainLinks() {
+    const links = await this.tag('chain_links')
+    return links !== undefined && Number(links) > 0
+  }
+
+  get hasHaplotypeIndex() {
+    return (
+      this.index.has('HaplotypeSamples') && this.index.has('HaplotypeLengths')
+    )
+  }
+
+  async haplotypeSampleInterval() {
+    const value = this.indexTags
+      ? this.indexTags.get('haplotype_index_interval')
+      : await this.tag('haplotype_index_interval')
+    return value === undefined ? undefined : Number(value)
+  }
+
+  private sampleFromRow(values: SqlValue[]): HaplotypeSample {
+    return {
+      node: num(values[0], 'HaplotypeSamples.node_handle'),
+      offset: num(values[1], 'HaplotypeSamples.node_offset'),
+      pathHandle: num(values[2], 'HaplotypeSamples.path_handle'),
+      orientation:
+        num(values[3], 'HaplotypeSamples.orientation') === 0
+          ? 'forward'
+          : 'reverse',
+      pathOffset: num(values[4], 'HaplotypeSamples.path_offset'),
+    }
+  }
+
+  async haplotypeSamplesInRange(minHandle: number, maxHandle: number) {
+    const samples: HaplotypeSample[] = []
+    for await (const key of this.index.indexScanFrom('HaplotypeSamples', [
+      minHandle,
+      0,
+    ])) {
+      const node = num(key[0], 'HaplotypeSamples.node_handle')
+      if (node > maxHandle) {
+        break
+      }
+      const row = await this.index.byRowid(
+        'HaplotypeSamples',
+        num(key[2], 'HaplotypeSamples rowid'),
+      )
+      if (row) {
+        samples.push(this.sampleFromRow(row))
+      }
+    }
+    return samples
+  }
+
+  async haplotypeSampleAt(node: number, offset: number) {
+    const key = await this.index.indexSeekLE('HaplotypeSamples', [node, offset])
+    if (key?.[0] !== node || key[1] !== offset) {
+      return undefined
+    }
+    const row = await this.index.byRowid(
+      'HaplotypeSamples',
+      num(key[2], 'HaplotypeSamples rowid'),
+    )
+    return row ? this.sampleFromRow(row) : undefined
+  }
+
+  async haplotypeLength(pathHandle: number) {
+    const row = await this.index.byRowid('HaplotypeLengths', pathHandle)
+    return row ? num(row[1], 'HaplotypeLengths.length') : undefined
+  }
+
+  async indexedPosition(
+    pathHandle: number,
+    pathOffset: number,
+  ): Promise<{ pathOffset: number; pos: Pos } | undefined> {
+    const key = await this.sqlite.indexSeekLE('ReferenceIndex', [
+      pathHandle,
+      pathOffset,
+    ])
+    if (key?.[0] !== pathHandle) {
+      return undefined
+    }
+    const row = await this.sqlite.byRowid(
+      'ReferenceIndex',
+      num(key[2], 'ReferenceIndex rowid'),
+    )
+    if (!row) {
+      throw new Error('ReferenceIndex row referenced by its index is missing')
+    }
+    return {
+      pathOffset: num(row[1], 'ReferenceIndex.path_offset'),
+      pos: {
+        node: num(row[2], 'ReferenceIndex.node_handle'),
+        offset: num(row[3], 'ReferenceIndex.node_offset'),
+      },
+    }
+  }
+}

package/src/filehandle.ts

@@ -0,0 +1,4 @@
+export interface ByteSource {
+  read(length: number, position: number): Promise<Uint8Array>
+  stat(): Promise<{ size: number }>
+}

package/src/gbwt/bytecode.ts

@@ -0,0 +1,87 @@
+export class ByteCodeReader {
+  offset = 0
+  private bytes: Uint8Array
+
+  constructor(bytes: Uint8Array) {
+    this.bytes = bytes
+  }
+
+  get done() {
+    return this.offset >= this.bytes.length
+  }
+
+  byte() {
+    const value = this.bytes[this.offset]
+    if (value === undefined) {
+      return undefined
+    }
+    this.offset += 1
+    return value
+  }
+
+  int() {
+    let shift = 1
+    let result = 0
+    while (this.offset < this.bytes.length) {
+      const value = this.bytes[this.offset]!
+      this.offset += 1
+      result += (value & 0x7f) * shift
+      shift *= 128
+      if ((value & 0x80) === 0) {
+        return result
+      }
+    }
+    return undefined
+  }
+}
+
+export interface Run {
+  value: number
+  len: number
+}
+
+const RLE_THRESHOLD = 255
+const RLE_UNIVERSE = 256
+
+export class RunReader {
+  private source: ByteCodeReader
+  private sigma: number
+  private threshold: number
+
+  constructor(bytes: Uint8Array, sigma: number) {
+    this.source = new ByteCodeReader(bytes)
+    this.sigma = sigma === 0 ? Number.MAX_SAFE_INTEGER : sigma
+    this.threshold =
+      this.sigma < RLE_THRESHOLD ? Math.floor(RLE_UNIVERSE / this.sigma) : 0
+  }
+
+  next(): Run | undefined {
+    if (this.sigma >= RLE_THRESHOLD) {
+      const value = this.source.int()
+      const len = this.source.int()
+      return value === undefined || len === undefined
+        ? undefined
+        : { value, len: len + 1 }
+    }
+    const byte = this.source.byte()
+    if (byte === undefined) {
+      return undefined
+    }
+    const value = byte % this.sigma
+    let len = Math.floor(byte / this.sigma) + 1
+    if (len === this.threshold) {
+      const extra = this.source.int()
+      if (extra === undefined) {
+        return undefined
+      }
+      len += extra
+    }
+    return { value, len }
+  }
+
+  *[Symbol.iterator]() {
+    for (let run = this.next(); run !== undefined; run = this.next()) {
+      yield run
+    }
+  }
+}

package/src/gbwt/node.ts

@@ -0,0 +1,64 @@
+export const ENDMARKER = 0
+
+export type Orientation = 'forward' | 'reverse'
+
+export function encodeNode(id: number, orientation: Orientation) {
+  return 2 * id + (orientation === 'reverse' ? 1 : 0)
+}
+
+export function nodeId(handle: number) {
+  return Math.floor(handle / 2)
+}
+
+export function nodeOrientation(handle: number): Orientation {
+  return handle % 2 === 0 ? 'forward' : 'reverse'
+}
+
+export function isReverse(handle: number) {
+  return handle % 2 === 1
+}
+
+export function flipNode(handle: number) {
+  return handle % 2 === 0 ? handle + 1 : handle - 1
+}
+
+export type NodeSide = 'left' | 'right'
+
+export function flipSide(side: NodeSide): NodeSide {
+  return side === 'left' ? 'right' : 'left'
+}
+
+export function entrySide(orientation: Orientation): NodeSide {
+  return orientation === 'forward' ? 'left' : 'right'
+}
+
+export function exitSide(orientation: Orientation): NodeSide {
+  return orientation === 'forward' ? 'right' : 'left'
+}
+
+export function entryOrientation(side: NodeSide): Orientation {
+  return side === 'left' ? 'forward' : 'reverse'
+}
+
+export function exitOrientation(side: NodeSide): Orientation {
+  return side === 'right' ? 'forward' : 'reverse'
+}
+
+export function edgeIsCanonical(from: number, to: number) {
+  const fromId = nodeId(from)
+  const toId = nodeId(to)
+  return isReverse(from)
+    ? toId > fromId || (toId === fromId && !isReverse(to))
+    : toId >= fromId
+}
+
+export function pathIsCanonical(path: number[]) {
+  const first = path[0]
+  const last = path[path.length - 1]
+  if (first === undefined || last === undefined) {
+    return true
+  }
+  return isReverse(first) === isReverse(last)
+    ? !isReverse(first)
+    : edgeIsCanonical(first, last)
+}

package/src/gbwt/record.ts

@@ -0,0 +1,177 @@
+import { ByteCodeReader, RunReader } from './bytecode.ts'
+import { ENDMARKER, flipNode, nodeId } from './node.ts'
+
+export interface Pos {
+  node: number
+  offset: number
+}
+
+export function decompressEdges(bytes: Uint8Array): Pos[] {
+  const reader = new ByteCodeReader(bytes)
+  const sigma = reader.int()
+  if (sigma === undefined || sigma === 0) {
+    return []
+  }
+  const edges: Pos[] = []
+  let prev = 0
+  for (let i = 0; i < sigma; i++) {
+    const delta = reader.int()
+    const offset = reader.int()
+    if (delta === undefined || offset === undefined) {
+      throw new Error('GBWT edge list ends early')
+    }
+    prev += delta
+    edges.push({ node: prev, offset })
+  }
+  return edges
+}
+
+export class GbwtRecord {
+  readonly edges: Pos[]
+  readonly bwt: Uint8Array
+
+  constructor(edges: Pos[], bwt: Uint8Array) {
+    this.edges = edges
+    this.bwt = bwt
+  }
+
+  runs() {
+    return new RunReader(this.bwt, this.edges.length)
+  }
+
+  private edge(rank: number) {
+    const edge = this.edges[rank]
+    if (edge === undefined) {
+      throw new Error(
+        `GBWT run refers to edge rank ${rank} of ${this.edges.length}`,
+      )
+    }
+    return edge
+  }
+
+  lf(i: number): Pos | undefined {
+    const offsets = this.edges.map(e => e.offset)
+    let offset = 0
+    for (const run of this.runs()) {
+      const edge = this.edge(run.value)
+      const soFar = offsets[run.value]!
+      if (offset + run.len > i) {
+        return edge.node === ENDMARKER
+          ? undefined
+          : { node: edge.node, offset: soFar + (i - offset) }
+      }
+      offsets[run.value] = soFar + run.len
+      offset += run.len
+    }
+    return undefined
+  }
+
+  private edgeTo(node: number) {
+    let low = 0
+    let high = this.edges.length
+    while (low < high) {
+      const mid = (low + high) >> 1
+      const edge = this.edges[mid]!
+      if (edge.node === node) {
+        return mid
+      }
+      if (edge.node < node) {
+        low = mid + 1
+      } else {
+        high = mid
+      }
+    }
+    return undefined
+  }
+
+  offsetTo(pos: Pos): number | undefined {
+    if (pos.node === ENDMARKER) {
+      return undefined
+    }
+    const rank = this.edgeTo(pos.node)
+    if (rank === undefined) {
+      return undefined
+    }
+    let succRank = this.edges[rank]!.offset
+    if (succRank > pos.offset) {
+      return undefined
+    }
+    let offset = 0
+    for (const run of this.runs()) {
+      offset += run.len
+      if (run.value !== rank) {
+        continue
+      }
+      succRank += run.len
+      if (succRank > pos.offset) {
+        return offset - (succRank - pos.offset)
+      }
+    }
+    return undefined
+  }
+
+  predecessorAt(i: number): number | undefined {
+    const counts = this.edges.map(e => ({
+      node: e.node === ENDMARKER ? ENDMARKER : flipNode(e.node),
+      count: 0,
+    }))
+    for (const run of this.runs()) {
+      const entry = counts[run.value]
+      if (entry) {
+        entry.count += run.len
+      }
+    }
+    for (let rank = 1; rank < counts.length; rank++) {
+      const prev = counts[rank - 1] as { node: number; count: number }
+      const curr = counts[rank] as { node: number; count: number }
+      if (nodeId(prev.node) === nodeId(curr.node)) {
+        counts[rank - 1] = curr
+        counts[rank] = prev
+      }
+    }
+    let offset = 0
+    for (const entry of counts) {
+      offset += entry.count
+      if (offset > i) {
+        return entry.node === ENDMARKER ? undefined : entry.node
+      }
+    }
+    return undefined
+  }
+
+  decompressArrays() {
+    const offsets = this.edges.map(e => e.offset)
+    let total = 0
+    for (const run of this.runs()) {
+      total += run.len
+    }
+    const nodes = new Int32Array(total)
+    const nextOffsets = new Int32Array(total)
+    let i = 0
+    for (const run of this.runs()) {
+      const edge = this.edge(run.value)
+      let offset = offsets[run.value]!
+      for (let k = 0; k < run.len; k++) {
+        nodes[i] = edge.node
+        nextOffsets[i] = offset
+        offset += 1
+        i += 1
+      }
+      offsets[run.value] = offset
+    }
+    return { nodes, offsets: nextOffsets }
+  }
+
+  decompress(): Pos[] {
+    const offsets = this.edges.map(e => e.offset)
+    const result: Pos[] = []
+    for (const run of this.runs()) {
+      const edge = this.edge(run.value)
+      for (let k = 0; k < run.len; k++) {
+        result.push({ node: edge.node, offset: offsets[run.value]! })
+        offsets[run.value] = offsets[run.value]! + 1
+      }
+    }
+    return result
+  }
+}