(* Copyright (C) 1997-1999 NEC Research Institute.
 * Please see the file LICENSE for license information.
 *)
structure Array2: ARRAY2 =
   struct
      open Primitive.Int
	 
      datatype 'a array =
	 Array of {rows: int,
		   cols: int,
		   array: 'a Array.array}

      fun dimensions(Array{rows, cols, ...}) = (rows, cols)
      fun nRows(Array{rows, ...}) = rows
      fun nCols(Array{cols, ...}) = cols

      type 'a region = {base: 'a array,
			row: int,
			col: int,
			nrows: int option,
			ncols: int option}

      fun checkRegion{base, row, col, nrows, ncols} =
	 let val (rows, cols) = dimensions base
	 in {stopRow = Array.checkSliceMax(row, nrows, rows),
	     stopCol = Array.checkSliceMax(col, ncols, cols)}
	 end
      
      fun wholeRegion a =
	 {base = a, row = 0, col = 0, nrows = NONE, ncols = NONE}

      datatype traversal = RowMajor | ColMajor

      local
	 fun make(rows, cols, doit) =
	    if safe andalso (rows < 0 orelse cols < 0)
	       then raise Size
	    else Array{rows = rows,
		       cols = cols,
		       array = doit(rows * cols)}
      in
	 fun arrayUninit(rows, cols) =
	    make(rows, cols, Primitive.Array.array)
	 fun array(rows, cols, init) =
	    make(rows, cols, fn size => Array.array(size, init))
      end

      fun array0() = Array{rows = 0,
			   cols = 0,
			   array = Primitive.Array.array 0}

      fun spot(Array{rows, cols, ...}, r, c) =
	 if safe andalso (geu(r, rows) orelse geu(c, cols))
	    then raise Subscript
	 else r * cols + c
	 
      fun sub(a as Array{array, ...}, r, c) =
	 Array.sub(array, spot(a, r, c))

      fun update(a as Array{array, ...}, r, c, x) =
	 Array.update(array, spot(a, r, c), x)

      fun 'a fromList(rows: 'a list list): 'a array =
	 case rows of
	    [] => array0()
	  | row1 :: _ =>
	       let
		  val cols = length row1
		  val a as Array{array, ...} = arrayUninit(length rows, cols)
	       in List.foldl (fn (row: 'a list, i) =>
			      let
				 val max = i + cols
				 val i' =
				    List.foldl (fn (x: 'a, i) =>
						(if i = max
						    then raise Size
						 else (Array.update(array, i, x);
						       i + 1)))
				    i row
			      in if i' = max
				    then i'
				 else raise Size
			      end)
		  0 rows;
		  a
	       end

      fun row(Array{rows, cols, array}, r) =
	 if safe andalso geu(r, rows)
	    then raise Subscript
	 else Array.extract(array, r * cols, SOME cols)

      fun column(a as Array{rows, cols, ...}, c) =
	 if safe andalso geu(c, cols)
	    then raise Subscript
	 else
           Vector.tabulate(rows, fn r => sub(a, r, c))

      fun foldi trv f b (region as {base, row, col, ...}) =
	 let
	    val {stopRow, stopCol} = checkRegion region
	 in
	    case trv of
	       RowMajor =>
		  Util.naturalFoldStartStop
		  (row, stopRow, b, fn (r, b) =>
		   Util.naturalFoldStartStop
		   (col, stopCol, b, fn (c, b) =>
		    f(r, c, sub(base, r, c), b)))
	     | ColMajor =>
		  Util.naturalFoldStartStop
		  (col, stopCol, b, fn (c, b) =>
		   Util.naturalFoldStartStop
		   (row, stopRow, b, fn (r, b) =>
		    f(r, c, sub(base, r, c), b)))
	 end

      fun fold trv f b a =
	  foldi trv (fn (_, _, x, b) => f(x, b)) b (wholeRegion a)

      fun appi trv f =
	 foldi trv (fn (r, c, x, ()) => f(r, c, x)) ()

      fun app trv f = fold trv (f o #1) ()

      fun modifyi trv f (r as {base, ...}) =
	 appi trv (fn (r, c, x) => update(base, r, c, f(r, c, x))) r

      fun modify trv f a = modifyi trv (f o #3) (wholeRegion a)

      fun tabulate trv (rows, cols, f) =
	 let val a = arrayUninit(rows, cols)
	 in modifyi trv (fn (r, c, _) => f(r, c)) (wholeRegion a);
	    a
	 end
   end