(* Lecture 18: Streams, Memoization *)
(* Author: Frank Pfenning *)

use "/afs/andrew/scs/cs/15-212-X/code/stream.sml";
use "/afs/andrew/scs/cs/15-212-X/code/dict.sml";

(* val stringMatch : string -> string -> int stream *)
fun stringMatch s p =
    let
      val lp = size (p)
      val ls = size (s)
      fun sm (i) = Stream.delay (fn () => sm' (i))
      and sm' (i) = if i+lp > ls then Stream.Empty
		    else if String.substring (s, i, lp) = p
			   then Stream.Cons (i, sm (i+1))
			 else sm'(i+1)
    in
      sm (0)
    end

fun append (s1,s2) = Stream.delay (fn () => append' (Stream.expose s1, s2))
and append' (Stream.Empty, s2) = Stream.expose s2
  | append' (Stream.Cons(x,s1), s2) =
      Stream.Cons(x, append (s1, s2))

structure IntPair :> ORDER where type t = int * int =
struct
  type t = int * int
  fun compare ((i1,j1),(i2,j2)) =
      (case Int.compare(i1,i2)
	 of EQUAL => Int.compare(j1,j2)
          | LESS => LESS
	  | GREATER => GREATER)
end;

signature BOARD = DICT where type key = int * int;

structure Board :> BOARD =
  RedBlackTree (structure Order = IntPair);

signature KNIGHTS_TOUR =
sig
  val solutions : int -> ((int * int) list) Stream.stream
end;

functor KnightsTour (structure Board : BOARD) :> KNIGHTS_TOUR =
struct

  datatype square = Clear | Visited
  fun initialize (b,n,(i,j)) =
        initialize' (Board.insert (b,((i,j),Clear)),n,(i,j))
  and initialize' (b,n,(1,1)) = b
    | initialize' (b,n,(i,1)) = initialize (b,n,(i-1,n))
    | initialize' (b,n,(i,j)) = initialize (b,n,(i,j-1))

  fun search (b,d,(i,j)) = Stream.delay (fn () => search' (b,d,(i,j)))
  and search' (b,0,(i,j)) = Stream.Cons ([(i,j)], Stream.empty)
    | search' (b,d,(i,j)) =
      let
	fun try ((di,dj),s) =
	       append (Stream.delay (fn () => move (b,d,(i+di,j+dj))), s)
      in
      Stream.expose
        (Stream.map (fn (l) => (i,j)::l)
	(List.foldl try Stream.empty
	 [(2,1),(1,2),(~1,2),(~2,1),(~2,~1),(~1,~2),(1,~2),(2,~1)]))
      end
	  
  and move (b,d,(i,j)) =
       (case Board.lookup b (i,j)
	  of NONE => Stream.Empty
           | SOME(Visited) => Stream.Empty
           | SOME(Clear) => search' (Board.insert(b,((i,j),Visited)),d-1,(i,j)))

  fun solutions n =
      search (initialize(Board.empty, n, (n,n)),
	      n*n, (1,1))
end;  (* functor KnightsTour *)

structure KnightsTour = KnightsTour (structure Board = Board);

(* Note that this implementation is NOT semantically *)
(* equivalent to the plain (non-memoizing) streams, since *)
(* effects will be executed only once in this implementation *)
structure BasicMemoStream :> BASIC_STREAM =
struct

  datatype 'a stream = Stream of unit -> 'a front
  and 'a front = Empty | Cons of 'a * 'a stream

  exception Uninitialized

  fun expose (Stream (d)) = d ()
  fun delay (d) =
      let val memo = ref (fn () => raise Uninitialized)
	  fun memoFun () =
	      let
		  val r = d ()
	       in
		  ( memo := (fn () => r) ; r )
	      end
	      handle exn => ( memo := (fn () => raise exn) ; raise exn )
      in
	memo := memoFun ; Stream (fn () => !memo ())
      end

  val empty = Stream (fn () => Empty)
  fun cons (x, s) = Stream (fn () => Cons (x, s))
end;

structure MStream :> STREAM =
  Stream (structure BasicStream = BasicMemoStream);

(* notMultiple p q >=> true iff q is not a multiple of p *)
fun notMultiple p q = (q mod p <> 0);

fun sieve s = MStream.delay (fn () => sieve' (MStream.expose s))
and sieve' (MStream.Empty) = MStream.Empty
  | sieve' (MStream.Cons(p, s)) =
      MStream.Cons (p, sieve (MStream.filter (notMultiple p) s));

val primes = sieve (MStream.tabulate (fn i => i+2));

signature MSTREAM_IO =
sig

  val readTerminal : string -> char MStream.stream
  val readFile : string -> char MStream.stream

end;

structure MStreamIO : MSTREAM_IO =
struct

  fun fromLine f nil = MStream.delay f
    | fromLine f (c::cs) = MStream.cons (c, fromLine f cs)

  fun readTerminal (prompt) =
      let
	fun getLine () =
	       ( TextIO.output (TextIO.stdOut, prompt) ;
		 TextIO.flushOut (TextIO.stdOut) ;
		 if TextIO.endOfStream TextIO.stdIn
		   then MStream.Empty
		 else
		   MStream.expose
		   (fromLine getLine 
		    (String.explode (TextIO.inputLine TextIO.stdIn))) )
      in
	MStream.delay getLine
      end

  fun readFile (fileName) =
      let
	val instream = TextIO.openIn fileName
        fun getLine () =
               if TextIO.endOfStream instream
		 then ( TextIO.closeIn instream ; MStream.Empty )
	       else
                 MStream.expose
		 (fromLine getLine
		  (String.explode (TextIO.inputLine instream)))
      in
	 MStream.delay getLine
      end
end;  (* structure MStreamIO *)