// add map / flatMap / filter to Stack
// ==> can use  for comprehension
//
// try:
// val s = EmptyStack push 1 push 2 push 3
// for ( x <- s ) yield (x * x)
//

abstract class Stack[+A] {
  def push[B >: A](x: B): Stack[B] = new NonEmptyStack(x, this)
  def isEmpty: Boolean
  def top: A
  def pop: Stack[A]

  def filter(p: A => Boolean): Stack[A] = {
    if (isEmpty) EmptyStack 
    else if (p( top )) new NonEmptyStack( top, pop.filter(p))
    else pop filter p
  }

  def map[Z](f: A => Z): Stack[Z] = {
    if (isEmpty) EmptyStack
    else new NonEmptyStack[Z]( f(top), pop.map(f))
  }

  def flatMap[Z](f: A => Stack[Z]): Stack[Z] = { 
    if (isEmpty) EmptyStack
    else { val first = f(top) ;
	   val rest = pop.flatMap(f) ; 
	  first.:::(rest)
	}
  }
   
  def :::[B >: A](that: Stack[B]): Stack[B] = {
    if (isEmpty) that
    else (pop.:::(that)).push(top)
  }

}

object EmptyStack extends Stack[Nothing] {
  def isEmpty = true
  def top = error("EmptyStack.top")
  def pop = error("EmptyStack.pop")

  override def toString = "Empty"
}

class NonEmptyStack[+A](elem: A, rest: Stack[A]) extends Stack[A] {
  def isEmpty = false
  def top = elem
  def pop = rest

  override def toString = "" + elem + " :: " + rest
}