// Graph-local SCF/memref inputs for `loom-lower-graph-memory`.
// Every sampled graph exports one memory capability in the `memories`
// segment of `dataflow.graph.return`: an imported graph memory input, a
// `memref.cast` view of that input, or a fresh `memref.alloc` root at the
// graph frontier. Bodies use normalized scalar `memref.load`/`memref.store`
// leaves over a canonical linear memory space, optionally nested in
// `scf.if` or source-sequential `scf.for`.

start: {new NG = random.randint(1, 3); new G = 0}
       'module {\n' graphs '}\n';

graphs: (G < NG) graph_def {G += 1} graphs
      | (G == NG) '';

graph_def: {new KIND = random.choice(['imported', 'view', 'fresh']); new SHAPE = random.choice(['plain', 'cond', 'loop'])}
           graph_pick;

graph_pick: (KIND == 'imported') imported_graph
          | (KIND == 'view') view_graph
          | (KIND == 'fresh') fresh_graph;

gid: [str(G)];

// An imported external memref capability is exported unchanged.
imported_graph: {new MEM = '%m'; new TY = 'memref<?xi32>'}
  '  dataflow.graph private @imported_export_' gid '(\n'
  '      %start: none, %lb: i64, %ub: i64, %step: i64, %i: index, %c: i1, %v: i32,\n'
  '      %m: memref<?xi32>) -> (memref<?xi32>)\n'
  '      attributes {input_segments = array<i32: 6, 0, 1>,\n'
  '                  result_segments = array<i32: 0, 0, 1>} {\n'
  body
  '    dataflow.graph.return values() streams()\n'
  '        memories(%m : memref<?xi32>) complete(%start : none)\n'
  '  }\n';

// A side-effect-free `memref.cast` view preserves the imported root.
view_graph: {new MEM = '%view'; new TY = 'memref<?xi32>'}
  '  dataflow.graph private @view_export_' gid '(\n'
  '      %start: none, %lb: i64, %ub: i64, %step: i64, %i: index, %c: i1, %v: i32,\n'
  '      %m: memref<4xi32>) -> (memref<?xi32>)\n'
  '      attributes {input_segments = array<i32: 6, 0, 1>,\n'
  '                  result_segments = array<i32: 0, 0, 1>} {\n'
  '    %view = memref.cast %m : memref<4xi32> to memref<?xi32>\n'
  body
  '    dataflow.graph.return values() streams()\n'
  '        memories(%view : memref<?xi32>) complete(%start : none)\n'
  '  }\n';

// A fresh `memref.alloc` result at the graph frontier is a unique
// invocation-local root exposed as an export.
fresh_graph: {new MEM = '%slot'; new TY = 'memref<4xi32>'}
  '  dataflow.graph private @fresh_export_' gid '(\n'
  '      %start: none, %lb: i64, %ub: i64, %step: i64, %i: index, %c: i1, %v: i32)\n'
  '      -> (memref<4xi32>)\n'
  '      attributes {input_segments = array<i32: 6, 0, 0>,\n'
  '                  result_segments = array<i32: 0, 0, 1>} {\n'
  '    %slot = memref.alloc() : memref<4xi32>\n'
  body
  '    dataflow.graph.return values() streams()\n'
  '        memories(%slot : memref<4xi32>) complete(%start : none)\n'
  '  }\n';

body: (SHAPE == 'plain') body_plain
    | (SHAPE == 'cond') body_cond
    | (SHAPE == 'loop') body_loop;

body_plain:
  '    memref.store %v, ' mem '[%i] : ' ty '\n'
  '    %loaded = memref.load ' mem '[%i] : ' ty '\n';

body_cond:
  '    scf.if %c {\n'
  '      memref.store %v, ' mem '[%i] : ' ty '\n'
  '    }\n'
  '    %loaded = memref.load ' mem '[%i] : ' ty '\n';

body_loop:
  '    scf.for %iv = %lb to %ub step %step : i64 {\n'
  '      %idx = arith.index_cast %iv : i64 to index\n'
  '      memref.store %v, ' mem '[%idx] : ' ty '\n'
  '    }\n'
  '    %loaded = memref.load ' mem '[%i] : ' ty '\n';

mem: [MEM];
ty: [TY];