// DFAs and NFAs based on Scala's partial functions

import scala.util.Try

// type abbreviation for partial functions

type :=>[A, B] = PartialFunction[A, B]

// some states for test cases 

abstract class State

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

// class for DFAs

case class DFA[A, C](start: A,                // starting state

                     delta: (A, C) :=> A,     // transition

                     fins:  A => Boolean) {   // final states

  // given a state and a "string", what is the next 

  // state, if there is any?

  def deltas(q: A, s: List[C]) : A = s match {

    case Nil => q

    case c::cs => deltas(delta(q, c), cs)

  }

  // is a "string" accepted by an DFA?

  def accepts(s: List[C]) : Boolean = 

    Try(fins(deltas(start, s))) getOrElse false

}

// DFA 1

val dtrans1 : (State, Char) :=> State = 

  { case (Q0, 'a') => Q0 

    case (Q0, 'b') => Q1 

  }

val dfa1 = DFA(Q0, dtrans1, Set[State](Q1))

dfa1.accepts("aaab".toList)     // true

dfa1.accepts("aacb".toList)     // false

// another DFA test

abstract class S

case object S0 extends S

case object S1 extends S

case object S2 extends S

case object Sink extends S

// transition function with a sink state

// S0 --a--> S1 --a--> S2

val sigma : (S, Char) :=> S = 

  { case (S0, 'a') => S1

    case (S1, 'a') => S2

    case _ => Sink

  }

val dfa1a = DFA(S0, sigma, Set[S](S2))

dfa1a.accepts("aa".toList)               //true

dfa1a.accepts("".toList)                 //false

dfa1a.accepts("ab".toList)               //false

// class for NFAs

case class NFA[A, C](starts: Set[A],            // starting states

                     delta: Set[(A, C) :=> A],  // transitions

                     fins:  A => Boolean) {     // final states 

  // given a state and a character, what is the set of next states?

  // if there is none => empty set

  def next(q: A, c: C) : Set[A] = 

  def nexts(qs: Set[A], c: C) : Set[A] =

    qs.flatMap(next(_, c))

  // given some states and a string, what is the set of next states?

  def deltas(qs: Set[A], s: List[C]) : Set[A] = s match {

    case Nil => qs

    case c::cs => deltas(nexts(qs, c), cs)

  }

  // is a string accepted by an NFA?

  def accepts(s: List[C]) : Boolean = 

    deltas(starts, s).exists(fins)

}

// NFA test cases

val trans1 = Set[(State, Char) :=> State](

  { 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(Set[State](Q0), trans1, Set[State](Q6))

nfa1.accepts("axaybzbc".toList)     // true

nfa1.accepts("aaaaxaybzbc".toList)  // true

nfa1.accepts("axaybzbd".toList)     // false

// 2

val trans2 = Set[(State, Char) :=> State](

  { case (Q0, 'a') => Q0 },

  { case (Q0, 'a') => Q1 },

  { case (Q0, 'b') => Q2 },

  { case (Q1, 'a') => Q1 },

  { case (Q2, 'b') => Q2 }

)

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

nfa2.accepts("aa".toList)             // false

nfa2.accepts("aaaaa".toList)          // false

nfa2.accepts("aaaaab".toList)         // true

nfa2.accepts("aaaaabbb".toList)       // true

nfa2.accepts("aaaaabbbaaa".toList)    // false

nfa2.accepts("ac".toList)             // false

// 3

val trans3 = Set[(State, Char) :=> State](

  { 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(Set[State](Q0), trans3, Set[State](Q5))

nfa3.accepts("aaaaabc".toList)      // true

nfa3.accepts("aaaabcd".toList)      // true

nfa3.accepts("aaaaab".toList)       // false

nfa3.accepts("aaaabc".toList)       // true

nfa3.accepts("aaaaabbbaaa".toList)  // false

// subset, or powerset, construction

def powerset[A, C](nfa: NFA[A, C]) : DFA[Set[A], C] = {

  DFA(nfa.starts, 

      { case (qs, c) => nfa.nexts(qs, c) },

      _.exists(nfa.fins))

}

val dfa2 = powerset(nfa1)

dfa2.accepts("axaybzbc".toList)     // true

dfa2.accepts("aaaaxaybzbc".toList)  // true

dfa2.accepts("axaybzbd".toList)     // false

val dfa3 = powerset(nfa2)

dfa3.accepts("aa".toList)             // false

dfa3.accepts("aaaaa".toList)          // false

dfa3.accepts("aaaaab".toList)         // true

dfa3.accepts("aaaaabbb".toList)       // true

dfa3.accepts("aaaaabbbaaa".toList)    // false

dfa3.accepts("ac".toList)             // false

// epsilon NFA

// fixpoint construction

import scala.annotation.tailrec

@tailrec

def fixpT[A](f: A => A, x: A): A = {

  val fx = f(x)

  if (fx == x) x else fixpT(f, fx) 

}

                      delta: Set[(A, Option[C]) :=> A],  // transition edges

                      fins:  A => Boolean) {             // final states 

  // epsilon transitions

  def enext(q: A) : Set[A] = 

    delta.flatMap((d) => Try(d(q, None)).toOption)

  def enexts(qs: Set[A]) : Set[A] = 

  // epsilon closure

  def ecl(qs: Set[A]) : Set[A] = 

    fixpT(enexts, qs)

  // "normal" transition

  def next(q: A, c: C) : Set[A] = 

    delta.flatMap((d) => Try(d(q, Some(c))).toOption)

  def nexts(qs: Set[A], c: C) : Set[A] = 

  def deltas(qs: Set[A], s: List[C]) : Set[A] = s match {

    case Nil => ecl(qs)

  }

  def accepts(s: List[C]) : Boolean = 

}

val etrans1 = Set[(State, Option[Char]) :=> State](

  { case (Q0, Some('a')) => Q1 },

  { case (Q1, None) => Q0 }

)

// Regular expressions fro derivative automata

abstract class Rexp

case object ZERO extends Rexp

case object ONE extends Rexp

case class CHAR(c: Char) extends Rexp 

case object ALL extends Rexp 

case object ALLPlus extends Rexp

case class ALT(r1: Rexp, r2: Rexp) extends Rexp 

case class SEQ(r1: Rexp, r2: Rexp) extends Rexp 

case class STAR(r: Rexp) extends Rexp 

case class NTIMES(r: Rexp, n: Int) extends Rexp 

case class UPNTIMES(r: Rexp, n: Int) extends Rexp 

case class NOT(r: Rexp) extends Rexp

case class AND(r1: Rexp, r2: Rexp) extends Rexp 

def get_chars(r: Rexp) : Set[Char] = r match {

  case ZERO => Set()

  case ONE => Set()

  case CHAR(c) => Set(c)

  case STAR(r) => get_chars(r)

  case NTIMES(r, _) => get_chars(r)

  case UPNTIMES(r, _) => get_chars(r)

  case NOT(r) => get_chars(r)

}

import scala.language.implicitConversions    

import scala.language.reflectiveCalls 

def charlist2rexp(s: List[Char]): Rexp = s match {

  case Nil => ONE

  case c::Nil => CHAR(c)

  case c::s => SEQ(CHAR(c), charlist2rexp(s))

}

implicit def string2rexp(s: String): Rexp = charlist2rexp(s.toList)

implicit def RexpOps (r: Rexp) = new {

  def | (s: Rexp) = ALT(r, s)

  def % = STAR(r)

  def ~ (s: Rexp) = SEQ(r, s)

}

implicit def stringOps (s: String) = new {

  def | (r: Rexp) = ALT(s, r)

  def | (r: String) = ALT(s, r)

  def % = STAR(s)

  def ~ (r: Rexp) = SEQ(s, r)

  def ~ (r: String) = SEQ(s, r)

}