func.decl import("input") @require(%condition: i32, %diagnostic: buffer)

func.decl import("json") @members(%json: buffer) -> (buffer, i32, i64)

func.decl import("prepare") @stage(%root: buffer, %tag: i64) -> (i32)

func.decl import("prepare") @config_i64(%stage: i32, %key: buffer, %value: i64)

func.decl import("prepare") @parameter(%stage: i32, %binding: i32, %input: i32, %source: i32, %policy: buffer)

func.decl import("prepare") @source_scope(%scope: buffer) -> (i32)

func.decl @tiny_storage(%maximum: i64, %span_capacity: i64, %context: i64, %pool: i64, %rows: i64, %pins: i64) -> (buffer, buffer, buffer, buffer, buffer)

// One shared command handles every host token/span bucket. The cold program
// owns storage geometry; the host owns physical page IDs and recurrent slots.
func.def public @text_prepare_v1(%prefill: i64, %context: i64, %pool: i64, %rows: i64, %pins: i64, %draft_depth: i32, %shapes: buffer, %shape_count: i64, %weights: buffer, %options: buffer) -> (i64, i32, buffer, buffer, buffer, buffer, buffer, buffer, buffer) {
  %c0_source = scalar.constant 0 : i32
  %origin = index.constant 0 : offset
  %c0 = index.constant 0 : index
  %c1 = index.constant 1 : index
  %c0_i32 = scalar.constant 0 : i32
  %c1_i32 = scalar.constant 1 : i32
  %c0_i64 = scalar.constant 0 : i64
  %c1_i64 = scalar.constant 1 : i64
  %c4_i64 = scalar.constant 4 : i64
  %pooled = scalar.cmpi ne, %pool, %c0_i64 : i64
  %packed = scalar.cmpi ne, %shape_count, %c0_i64 : i64
  %no_draft = scalar.cmpi eq, %draft_depth, %c0_i32 : i32
  %has_epochs = scalar.andi %pooled, %packed : i1
  %supported_mode = scalar.andi %has_epochs, %no_draft : i1
  %fields, %type, %field_count = func.call @members(%options) : (buffer) -> (buffer, i32, i64)
  %object_type = scalar.constant 2 : i32
  %object = scalar.cmpi eq, %type, %object_type : i32
  %empty_options = scalar.cmpi eq, %field_count, %c0_i64 : i64
  %valid_options = scalar.andi %object, %empty_options : i1
  %supported = scalar.andi %supported_mode, %valid_options : i1
  %supported32 = scalar.extui %supported : i1 to i32
  %error = buffer.constant utf8("text tiny requires pooled epochs without MTP") : buffer
  func.call @require(%supported32, %error) : (i32, buffer)
  %root = buffer.constant utf8("advance") : buffer
  %stage = func.call @stage(%root, %c0_i64) : (buffer, i64) -> (i32)
  %context_key = buffer.constant utf8("tiny.context") : buffer
  func.call @config_i64(%stage, %context_key, %context) : (i32, buffer, i64)
  %empty = buffer.constant utf8("") : buffer
  %policy = buffer.constant utf8("weights.loom") : buffer
  %scope = buffer.constant utf8("model") : buffer
  %checkpoint = func.call @source_scope(%scope) : (buffer) -> (i32)
  func.call @parameter(%stage, %c0_i32, %c0_source, %checkpoint, %policy) : (i32, i32, i32, i32, buffer)
  %state_bytes = index.constant 128 : offset
  %state = buffer.alloca<private> align(8) %state_bytes : buffer
  %state_view = buffer.view %state[%origin] : buffer -> view<1xi32>
  view.store %stage, %state_view[0] : i32, view<1xi32>

  %bounded_count = scalar.assume %shape_count [range(%shape_count, 1, 2147483647)] : i64
  %count = index.cast %bounded_count : i64 to index
  %shape_view = buffer.view %shapes[%origin] : buffer -> view<[%count]x2xi64>
  %maximum, %span_capacity = scf.for %i = [%c0 to %count step %c1](%previous = %c0_i64 : i64, %previous_spans = %c0_i64 : i64) -> (i64, i64) {
    %tokens = view.load %shape_view[%i, 0] : view<[%count]x2xi64> -> i64
    %larger = scalar.cmpi ugt, %tokens, %previous : i64
    %next = scf.select %larger, %tokens, %previous : i64
    %spans = view.load %shape_view[%i, 1] : view<[%count]x2xi64> -> i64
    %more_spans = scalar.cmpi ugt, %spans, %previous_spans : i64
    %next_spans = scf.select %more_spans, %spans, %previous_spans : i64
    scf.yield %next, %next_spans : i64, i64
  }
  %smaller = scalar.cmpi ult, %maximum, %prefill : i64
  %effective = scf.select %smaller, %maximum, %prefill : i64
  %tokens_key = buffer.constant utf8("tiny.tokens") : buffer
  %spans_key = buffer.constant utf8("tiny.spans") : buffer
  func.call @config_i64(%stage, %spans_key, %span_capacity) : (i32, buffer, i64)
  func.call @config_i64(%stage, %tokens_key, %maximum) : (i32, buffer, i64)
  %layout_origin = index.constant 88 : offset
  %layout = buffer.view %state[%layout_origin] : buffer -> view<5xi64>
  view.store %maximum, %layout[0] : i64, view<5xi64>
  view.store %span_capacity, %layout[1] : i64, view<5xi64>
  view.store %context, %layout[2] : i64, view<5xi64>
  view.store %pool, %layout[3] : i64, view<5xi64>
  view.store %pins, %layout[4] : i64, view<5xi64>
  %allocations, %views, %origins, %draft, %geometry = func.call @tiny_storage(%maximum, %span_capacity, %context, %pool, %rows, %pins) : (i64, i64, i64, i64, i64, i64) -> (buffer, buffer, buffer, buffer, buffer)
  %terminal = buffer.constant utf8("#") : buffer
  func.return %effective, %c1_i32, %state, %terminal, %allocations, %views, %origins, %empty, %geometry : i64, i32, buffer, buffer, buffer, buffer, buffer, buffer, buffer
}