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

// edges of the DFAs and NFAs

type Edge[A, C] = PartialFunction[(A, C), A]

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

                     trans: Edge[A, C],      // transition

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

  // given a state and a character, 

  // what is the next state, if there is any?

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

    trans.lift.apply((q, c))

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

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

    case Nil => Some(q)

    case c::cs => next(q, c).flatMap(nexts(_, cs))

  }

  // is a string accepted by an DFA?

  def accepts(s: List[C]) : Boolean = nexts(start, s) match {

    case None => false

    case Some(q) => fins(q)

  }

}

// DFA 1

val dtrans1 : Edge[State, Char] = 

  { 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

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

val sigma : Edge[S, Char] = 

  { 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

                     trans: Set[Edge[A, C]],    // transition edges

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

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

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

    trans.flatMap(_.lift.apply((q, c)))

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

  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 nextss(qs: Set[A], s: List[C]) : Set[A] = s match {

    case Nil => qs

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

  }

  // is a string accepted by an NFA?

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

    nextss(starts, s.toList).exists(fins)

}

// NFA test cases

// 1 

val trans1 = Set[Edge[State, Char]](

  { 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[Edge[State, Char]](

  { 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[Edge[State, Char]](

  { 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

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) 

}

case class eNFA[A, C](starts: Set[A],                    // starting state

                      trans: Set[Edge[A, Option[C]]],    // transition edges

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

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

    trans.flatMap(_.lift.apply((q, None)))

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

    qs ++ qs.flatMap(enext(_))

  // epsilon closure

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

    fixpT(enexts, qs)

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

    trans.flatMap(_.lift.apply((q, Some(c))))

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

    qs.flatMap(next(_, c))

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

    case Nil => ecl(qs)

    case c::cs => nextss(ecl(nexts(ecl(qs), c)), cs)

  }

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

    nextss(starts, s.toList).exists(fins)

}

val etrans1 = Set[Edge[State, Option[Char]]](

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

  { case (Q1, None) => Q0 }

)

val enfa = eNFA(Set[State](Q0), etrans1, Set[State](Q1))

enfa.accepts("a".toList)              // true

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

enfa.accepts("aaaaa".toList)          // true

enfa.accepts("aaaaab".toList)         // flase

enfa.accepts("aaaaabbb".toList)       // false

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

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

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

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)

}

// nullable function: tests whether the regular 

// expression can recognise the empty string

def nullable (r: Rexp) : Boolean = r match {

  case ZERO => false

  case ONE => true

  case CHAR(_) => false

  case ALL => false

  case ALT(r1, r2) => nullable(r1) || nullable(r2)

  case SEQ(r1, r2) => nullable(r1) && nullable(r2)

  case STAR(_) => true

  case NTIMES(r, i) => if (i == 0) true else nullable(r)

  case UPNTIMES(r: Rexp, n: Int) => true

}

// derivative of a regular expression w.r.t. a character

def der (c: Char, r: Rexp) : Rexp = r match {

  case ZERO => ZERO

  case ONE => ZERO

  case CHAR(d) => if (c == d) ONE else ZERO

  case ALL => ONE

  case ALT(r1, r2) => ALT(der(c, r1), der(c, r2))

  case SEQ(r1, r2) => 

    if (nullable(r1)) ALT(SEQ(der(c, r1), r2), der(c, r2))

    else SEQ(der(c, r1), r2)

  case STAR(r1) => SEQ(der(c, r1), STAR(r1))

  case NTIMES(r1, i) => 

    if (i == 0) ZERO else

    if (nullable(r1)) SEQ(der(c, r1), UPNTIMES(r1, i - 1))

    else SEQ(der(c, r1), NTIMES(r1, i - 1))

  case UPNTIMES(r1, i) => 

    if (i == 0) ZERO

    else SEQ(der(c, r1), UPNTIMES(r1, i - 1)) 

}

// partial derivative of a regular expression w.r.t. a character

def pder (c: Char, r: Rexp) : Set[Rexp] = r match {

  case ZERO => Set()

  case ONE => Set()

  case CHAR(d) => if (c == d) Set(ONE) else Set()

  case ALL => Set(ONE)

  case ALT(r1, r2) => pder(c, r1) ++ pder(c, r2)

  case SEQ(r1, r2) => 

    (for (pr1 <- pder(c, r1)) yield SEQ(pr1, r2)) ++ 

    (if (nullable(r1)) pder(c, r2) else Set())

  case STAR(r1) => 

    for (pr1 <- pder(c, r1)) yield SEQ(pr1, STAR(r1))

  case NTIMES(r1, i) => 

    if (i == 0) Set() else

    if (nullable(r1)) 

      for (pr1 <- pder(c, r1)) yield SEQ(pr1, UPNTIMES(r1, i - 1))

    else 

      for (pr1 <- pder(c, r1)) yield SEQ(pr1, NTIMES(r1, i - 1))

  case UPNTIMES(r1, i) => 

    if (i == 0) Set()

    else 

      for (pr1 <- pder(c, r1)) yield SEQ(pr1, UPNTIMES(r1, i - 1))

}

val r = STAR(ALL) ~ "a" ~ NTIMES(ALL, 3) ~ "bc"

val pder_nfa = NFA[Set[Rexp], Char](Set(Set(r)), 

                                    Set( { case (rs, c) => rs.flatMap(pder(c, _))} ), 

                                    _.exists(nullable))

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

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

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

///

type Trans[A, C] = Map[(A, C), A]