type Node = struct { edge_start: u32; num_edges: u32; };
type StopProd = struct { stop: bool; };

fn make_stop_prod() -> StopProd {
  let ret : StopProd;
  ret.stop = true;
  return ret;
}

#[entry]
fn bfs<n, m: usize>(graph_nodes: Node[n], source: u32, edges: u32[m]) -> i32[n] {
  let stop = false;

  // The mask selects the set of nodes that we consider in each iteration
  // It includes only the nodes that were visited for the first time in the
  // prior iteration (and in the first iteration just the source node)
  let mask: bool[n];
  mask[source as u64] = true;

  let visited: bool[n];
  visited[source as u64] = true;
  
  @cost @cost_init let cost: i32[n];
  @cost_init for i in 0..n {
    cost[i] = -1;
  }
  cost[source as u64] = 0;

  // Nodes that were updated in the current iteration
  let updated: bool[n];

  while !stop {
    @loop1 for i in 0..n {
      if mask[i] {
        mask[i] = false;

        let edge_start = graph_nodes[i].edge_start as u64;
        let num_edges = graph_nodes[i].num_edges as u64;

        for edge in edge_start..edge_start + num_edges {
          let id = edges[edge] as u64;
          if !visited[id] {
            cost[id] = cost[i] + 1;
            updated[id] = true;
          }
        }
      }
    }

    @make let stop_prod = make_stop_prod();
    @loop2 for i in 0..n {
      if updated[i] {
        mask[i] = true;
        visited[i] = true;
        updated[i] = false;
        stop_prod.stop = updated[i];
      }
    }
    stop = stop_prod.stop;
  }

  return cost;
}