// NFAs based on Scala's partial functions


// function that tests whether the intersection of two sets
// is non-empty
def share[A](a: Set[A], b: Set[A]) : Boolean =
  !(a intersect b).isEmpty

share(Set(1,2,3), Set(2, 3, 4))   // true
share(Set(1,2,3), Set(4, 5, 6))   // false

// nodes of the NFAs
abstract class State
type States = Set[State]

// edges of the NFAs
type Edge = PartialFunction[(State, Char), State]
type Edges = Set[Edge]


// gives a new state, when the edge fires
def fire(e: Edge, qc: (State, Char)) : Option[State] = 
  if (e.isDefinedAt(qc)) Some(e(qc)) else None


// class for NFAs
case class NFA(start: State,       // starting state
               trans: Edges,       // transition edges
               fins:  States)      // final states


// given the transitions and a state/character, 
// what are the set of next states?
def next(trans: Edges, q: State, c: Char) : States = {
  trans.map(fire(_, (q, c))).flatten
}

// given the transitions and some states/character, 
// what are the set of next states?
def nexts(trans: Edges, qs: States, c: Char) : States = {
  qs.flatMap(next(trans, _, c))
}


// given the transitions, some states and a string, 
// what are the set of next states?
def nextss(trans: Edges, qs: States, s: List[Char]) : States = s match {
  case Nil => qs
  case c::cs => nextss(trans, nexts(trans, qs, c), cs)
}

// is a string accepted by an NFA?
def accepts(nfa: NFA, s: String) : Boolean = {
  share(nextss(nfa.trans, Set(nfa.start), s.toList), nfa.fins)
}


// test cases
case object Q0 extends State
case object Q1 extends State
case object Q2 extends State
case object Q3 extends State
case object Q4 extends State
case object Q5 extends State
case object Q6 extends State

// 1 
val trans1 = Set[Edge](
  { case (Q0, 'a') => Q1 },
  { case (Q0, _)   => Q0 },
  { case (Q1, _)   => Q2 },
  { case (Q2, _)   => Q3 },
  { case (Q3, _)   => Q4 },
  { case (Q4, 'b') => Q5 },
  { case (Q5, 'c') => Q6 }
)

val nfa1 = NFA(Q0, trans1, Set[State](Q6))

accepts(nfa1, "axaybzbc")     // true
accepts(nfa1, "aaaaxaybzbc")  // true
accepts(nfa1, "axaybzbd")     // false


// 2
val trans2 = Set[Edge](
  { case (Q0, 'a') => Q0 },
  { case (Q0, 'a') => Q1 },
  { case (Q0, 'b') => Q2 },
  { case (Q1, 'a') => Q1 },
  { case (Q2, 'b') => Q2 }
)

val nfa2 = NFA(Q0, trans2, Set[State](Q2))

accepts(nfa2, "aa")             // false
accepts(nfa2, "aaaaa")          // false
accepts(nfa2, "aaaaab")         // true
accepts(nfa2, "aaaaabbb")       // true
accepts(nfa2, "aaaaabbbaaa")    // false
accepts(nfa2, "ac")             // false

// 3
val trans3 = Set[Edge](
  { case (Q0, _)   => Q0 },
  { case (Q0, 'a') => Q1 },
  { case (Q0, 'b') => Q3 },
  { case (Q1, 'b') => Q2 },
  { case (Q2, 'c') => Q5 },
  { case (Q3, 'c') => Q4 },
  { case (Q4, 'd') => Q5 }
)

val nfa3 = NFA(Q0, trans3, Set[State](Q5))

accepts(nfa3, "aaaaabc")      // true
accepts(nfa3, "aaaabcd")      // true
accepts(nfa3, "aaaaab")       // false
accepts(nfa3, "aaaabc")       // true
accepts(nfa3, "aaaaabbbaaa")  // false