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

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

  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)

  // depth-first search version of accept

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

    case Nil => fins(q)

    case c::cs => delta.exists(d => Try(search(d(q, c), cs)) getOrElse false)

  }

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

    starts.exists(search(_, s))

}

// NFA test cases

// 1:  NFA for STAR(ALL) ~ "a" ~ NTIMES(ALL, 3) ~ "bc"

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

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

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

nfa1.accepts2("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

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

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

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

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

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

nfa2.accepts2("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) 

}

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

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

    qs | qs.flatMap(enext(_))

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

    ecl(ecl(qs).flatMap(next(_, c)))

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

    case Nil => ecl(qs)

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

  }

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

    deltas(Set(start), s.toList).exists(fins)

}

// test cases for eNFAs

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

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

  { case (Q1, None) => Q0 }

)

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

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

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

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

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

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

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

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

// example from handouts 

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

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

  { case (Q0, None) => Q1 },

  { case (Q0, None) => Q2 },

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

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

  { case (Q1, None) => Q0 }

)

val enfa2 = eNFA(Q0, etrans2, Set[State](Q2))

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

enfa2.accepts("".toList)               // true

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

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

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

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

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

// states for Thompson construction

case class TState(i: Int) extends State

object TState {

  var counter = 0

  def apply() : TState = {

    counter += 1;

    new TState(counter - 1)

  }

}

// eNFA that does not accept any string

def eNFA_ZERO(): eNFA[TState, Char] = {

  val Q = TState()

  eNFA(Q, Set(), Set())

}

// eNFA that accepts the empty string

def eNFA_ONE() : eNFA[TState, Char] = {

  val Q = TState()

  eNFA(Q, Set(), Set(Q))

}

// eNFA that accepts the string "c"

def eNFA_CHAR(c: Char) : eNFA[TState, Char] = {

  val Q1 = TState()

  val Q2 = TState()

  eNFA(Q1, 

       Set({ case (Q1, Some(d)) if (c == d) => Q2 }),

       Set(Q2))

}

// alternative of two eNFAs

def eNFA_ALT(enfa1: eNFA[TState, Char], enfa2: eNFA[TState, Char]) : eNFA[TState, Char] = {

  val Q = TState()

  eNFA(Q,

       enfa1.delta | enfa2.delta |

       Set({ case (Q, None) => enfa1.start },

           { case (Q, None) => enfa2.start }),

       q => enfa1.fins(q) || enfa2.fins(q))

}

// sequence of two eNFAs

def eNFA_SEQ(enfa1: eNFA[TState, Char], enfa2: eNFA[TState, Char]) : eNFA[TState, Char] = {

  eNFA(enfa1.start,

       enfa1.delta | enfa2.delta | 

       Set({ case (q, None) if enfa1.fins(q) => enfa2.start }),

       enfa2.fins)

}

// star of an eNFA

def eNFA_STAR(enfa: eNFA[TState, Char]) : eNFA[TState, Char] = {

  val Q = TState()

  eNFA(Q,

       enfa.delta |

       Set({ case (q, None) if enfa.fins(q) => Q },

           { case (Q, None) => enfa.start }),

       Set(Q))

}

/*

def tofunset[A, C](d: (A, Option[C]) :=> Set[A])(q: A, c: C) : Set[(A, C) :=> A] = {

  d((q, Some(c))).map ((qs: A) => { case (qi, ci) if (qi == q && ci == c) => qs } : (A, C) :=> A)

}

def eNFA2NFA[A, C](starts: Set[A],                    // starting state

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

    qs | qs.flatMap(enext(_))

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

    ecl(ecl(qs).flatMap(next(_, c)))

  def nfa_delta : Set[(A, C) :=> A] = delta.flatMap(d => tofunset(d))

  def nfa_starts = ecl(starts)

  def nfa_fins = (q: A) => ecl(Set[A](q)) exists fins

  NFA(nfa_starts, nfa_delta, nfa_fins)

}

*/ 

// Regular expressions for 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 ALT(r1, r2) => get_chars(r1) | get_chars(r2)

  case SEQ(r1, r2) => get_chars(r1) | get_chars(r2)

  case STAR(r) => get_chars(r)

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

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

  case ALL => ('a' to 'z').toSet | Set('*','/','\\')

  case NOT(r) => get_chars(r)

  case AND(r1, r2) => get_chars(r1) | get_chars(r2)

}

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)

}

// thompson construction only for basic regular expressions

def thompson (r: Rexp) : eNFA[TState, Char] = r match {

  case ZERO => eNFA_ZERO()

  case ONE => eNFA_ONE()

  case CHAR(c) => eNFA_CHAR(c)  

  case ALT(r1, r2) => eNFA_ALT(thompson(r1), thompson(r2))

  case SEQ(r1, r2) => eNFA_SEQ(thompson(r1), thompson(r2))

  case STAR(r1) => eNFA_STAR(thompson(r1))

}

// regular expression matcher using Thompson's

def tmatcher(r: Rexp, s: String) : Boolean = thompson(r).accepts(s.toList)

// test cases for thompson construction

tmatcher(ZERO, "")   // false

tmatcher(ZERO, "a")  // false

tmatcher(ONE, "")    // true

tmatcher(ONE, "a")   // false

tmatcher(CHAR('a'), "")    // false

tmatcher(CHAR('a'), "a")   // true

tmatcher(CHAR('a'), "b")   // false

tmatcher("a" | "b", "")    // false

tmatcher("a" | "b", "a")   // true

tmatcher("a" | "b", "b")   // true

tmatcher("a" | "b", "c")   // false

tmatcher("a" | "b", "ab")  // false

tmatcher("a" ~ "b", "")    // false

tmatcher("a" ~ "b", "a")   // false

tmatcher("a" ~ "b", "b")   // false

tmatcher("a" ~ "b", "c")   // false

tmatcher("a" ~ "b", "ab")  // true

tmatcher("a" ~ "b", "aba") // false

tmatcher(EVIL2, "aaaaaab")   // true

tmatcher(EVIL2, "aaaaaa")    // false

tmatcher(EVIL2, "a" * 100)   // false

//optional

def OPT(r: Rexp) = ALT(r, ONE)

//n-times

def NTIMES(r: Rexp, n: Int) : Rexp = n match {

  case 0 => ONE

  case 1 => r

  case n => SEQ(r, NTIMES(r, n - 1))

}

// evil regular exproession

def EVIL(n: Int) = SEQ(NTIMES(OPT("a"), n), NTIMES("a", n))

val EVIL2 = STAR(STAR("a")) ~ "b"

// helper function for recording time

def time_needed[T](i: Int, code: => T) = {

  val start = System.nanoTime()

  for (j <- 1 to i) code

  val end = System.nanoTime()

  (end - start)/(i * 1.0e9)

}