// 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