// uses bitcode sequences and annotated regular 
// expressions
//
// for basic regular expressions and RECD
//
// uses a non-standard extraction method for generating
// tokens (this is tail-recursive) 
//
// can match 60 copies of the fib-program (size 10500)
// in about 20 secs


import scala.language.implicitConversions    
import scala.language.reflectiveCalls
import scala.annotation.tailrec   

// standard regular expressions
abstract class Rexp 
case object ZERO extends Rexp
case object ONE extends Rexp
case class CHAR(c: Char) 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 RECD(x: String, r: Rexp) extends Rexp

abstract class Bit
case object Z extends Bit
case object S extends Bit

type Bits = List[Bit]

// annotated regular expressions
abstract class ARexp 
case object AZERO extends ARexp
case class AONE(bs: Bits) extends ARexp
case class ACHAR(bs: Bits, c: Char) extends ARexp
case class AALTS(bs: Bits, rs: List[ARexp]) extends ARexp 
case class ASEQ(bs: Bits, r1: ARexp, r2: ARexp) extends ARexp 
case class ASTAR(bs: Bits, r: ARexp) extends ARexp 

// an abbreviation for binary alternatives
def AALT(bs: Bits, r1: ARexp, r2: ARexp) = AALTS(bs, List(r1, r2))

abstract class Val
case object Empty extends Val
case class Chr(c: Char) extends Val
case class Sequ(v1: Val, v2: Val) extends Val
case class Left(v: Val) extends Val
case class Right(v: Val) extends Val
case class Stars(vs: List[Val]) extends Val
case class Recd(x: String, v: Val) extends Val

abstract class BTree
case object LZ extends BTree
case class L(bs: Bits) extends BTree
case class LALTS(bs: Bits, bts: List[BTree]) extends BTree 
case class LSEQ(bs: Bits, bt1: BTree, bt2: BTree) extends BTree 
case class LSTAR(bs: Bits, bt: BTree) extends BTree

def ext(r: ARexp): BTree = r match {
  case AZERO => LZ
  case AONE(bs) => L(bs)
  case ACHAR(bs, _) => L(bs)
  case AALTS(bs, rs) => LALTS(bs, rs.map(ext))
  case ASEQ(bs, r1, r2) => LSEQ(bs, ext(r1), ext(r2))
  case ASTAR(bs, r) => LSTAR(bs, ext(r))
}

// annotated regular expressions
abstract class BRexp 
case object BZERO extends BRexp
case object BONE extends BRexp
case class BCHAR(c: Char) extends BRexp
case class BALTS(rs: List[BRexp]) extends BRexp 
case class BSEQ(r1: BRexp, r2: BRexp) extends BRexp 
case class BSTAR(r: BRexp) extends BRexp 

def ex(r: ARexp): BRexp = r match {
  case AZERO => BZERO
  case AONE(_) => BONE
  case ACHAR(_, c) => BCHAR(c)
  case AALTS(_, rs) => BALTS(rs.map(ex))
  case ASEQ(_, r1, r2) => BSEQ(ex(r1), ex(r2))
  case ASTAR(_, r) => BSTAR(ex(r))
}

   
// some convenience for typing in regular expressions
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)
  def $ (r: Rexp) = RECD(s, r)
}

def size(r: Rexp) : Int = r match {
  case ZERO => 1
  case ONE => 1
  case CHAR(_) => 1
  case ALT(r1, r2) => 1 + size(r1) + size(r2)
  case SEQ(r1, r2) => 1 + size(r1) + size(r2)
  case STAR(r) => 1 + size(r)
  case RECD(_, r) => 1 + size(r)
}



// Bitcoded + Annotation
//=======================

//erase function: extracts the Rexp from ARexp
def erase(r:ARexp): Rexp = r match{
  case AZERO => ZERO
  case AONE(_) => ONE
  case ACHAR(bs, c) => CHAR(c)
  case AALTS(bs, Nil) => ZERO
  case AALTS(bs, r::Nil) => erase(r)
  case AALTS(bs, r::rs) => ALT(erase(r), erase(AALTS(bs, rs)))
  case ASEQ(bs, r1, r2) => SEQ (erase(r1), erase(r2))
  case ASTAR(cs, r)=> STAR(erase(r))
}

def fuse(bs: Bits, r: ARexp) : ARexp = r match {
  case AZERO => AZERO
  case AONE(cs) => AONE(bs ++ cs)
  case ACHAR(cs, c) => ACHAR(bs ++ cs, c)
  case AALTS(cs, rs) => AALTS(bs ++ cs, rs)
  case ASEQ(cs, r1, r2) => ASEQ(bs ++ cs, r1, r2)
  case ASTAR(cs, r) => ASTAR(bs ++ cs, r)
}

def internalise(r: Rexp) : ARexp = r match {
  case ZERO => AZERO
  case ONE => AONE(Nil)
  case CHAR(c) => ACHAR(Nil, c)
  case ALT(r1, r2) => AALT(Nil, fuse(List(Z), internalise(r1)), fuse(List(S), internalise(r2)))
  case SEQ(r1, r2) => ASEQ(Nil, internalise(r1), internalise(r2))
  case STAR(r) => ASTAR(Nil, internalise(r))
  case RECD(_, r) => internalise(r)
}

// example
// internalise(("a" | "ab") ~ ("b" | ""))


// decoding of a value from a bitsequence
// (this is not tail-recursive and therefore a potential bottleneck)
def vdecode_aux(r: Rexp, bs: Bits) : (Val, Bits) = (r, bs) match {
  case (ONE, bs) => (Empty, bs)
  case (CHAR(c), bs) => (Chr(c), bs)
  case (ALT(r1, r2), Z::bs) => {
    val (v, bs1) = vdecode_aux(r1, bs)
    (Left(v), bs1)
  }
  case (ALT(r1, r2), S::bs) => {
    val (v, bs1) = vdecode_aux(r2, bs)
    (Right(v), bs1)
  }
  case (SEQ(r1, r2), bs) => {
    val (v1, bs1) = vdecode_aux(r1, bs)
    val (v2, bs2) = vdecode_aux(r2, bs1)
    (Sequ(v1, v2), bs2)
  }
  case (STAR(r1), Z::bs) => {
    val (v, bs1) = vdecode_aux(r1, bs)
    val (Stars(vs), bs2) = vdecode_aux(STAR(r1), bs1)
    (Stars(v::vs), bs2)
  }
  case (STAR(_), S::bs) => (Stars(Nil), bs)
  case (RECD(s, r1), bs) => 
    val (v, bs1) = vdecode_aux(r1, bs)
    (Recd(s, v), bs1)
}

def vdecode(r: Rexp, bs: Bits) = vdecode_aux(r, bs) match {
  case (v, Nil) => v
  case _ => throw new Exception("Not decodable")
}

// decoding of sequence of string tokens from a bitsequence
// tail-recursive version using an accumulator (alternative for 
// vdecode)
@tailrec
def sdecode_aux(rs: List[Rexp], bs: Bits, acc: List[String]) : List[String] = (rs, bs) match {
  case (Nil, _) => acc
  case (_, Nil) => acc
  case (ONE::rest, bs) => sdecode_aux(rest, bs, acc)
  case (CHAR(c)::rest, bs) => 
    sdecode_aux(rest, bs, (acc.head ++ c.toString)::acc.tail)
  case (ALT(r1, r2)::rest, Z::bs) => sdecode_aux(r1::rest, bs, acc)
  case (ALT(r1, r2)::rest, S::bs) => sdecode_aux(r2::rest, bs, acc)
  case (SEQ(r1, r2)::rest, bs) => sdecode_aux(r1::r2::rest, bs, acc)
  case (STAR(r1)::rest, Z::bs) => sdecode_aux(r1::STAR(r1)::rest, bs, acc)
  case (STAR(_)::rest, S::bs) => sdecode_aux(rest, bs, acc)
  case (RECD(s, r1)::rest, bs) => 
    sdecode_aux(r1::rest, bs, s"$s:"::acc)
}

def sdecode(r: Rexp, bs: Bits) : List[String] = 
  sdecode_aux(List(r), bs, List("")).reverse.tail



// nullable function: tests whether the an (annotated) 
// regular expression can recognise the empty string
def bnullable (r: ARexp) : Boolean = r match {
  case AZERO => false
  case AONE(_) => true
  case ACHAR(_,_) => false
  case AALTS(_, rs) => rs.exists(bnullable)
  case ASEQ(_, r1, r2) => bnullable(r1) && bnullable(r2)
  case ASTAR(_, _) => true
}

def bmkeps(r: ARexp) : Bits = r match {
  case AONE(bs) => bs
  case AALTS(bs, r::Nil) => bs ++ bmkeps(r) 
  case AALTS(bs, r::rs) => 
    if (bnullable(r)) bs ++ bmkeps(r) else bmkeps(AALTS(bs, rs))  
  case ASEQ(bs, r1, r2) => bs ++ bmkeps(r1) ++ bmkeps(r2)
  case ASTAR(bs, r) => bs ++ List(S)
}

// derivative of a regular expression w.r.t. a character
def bder(c: Char, r: ARexp) : ARexp = r match {
  case AZERO => AZERO
  case AONE(_) => AZERO
  case ACHAR(bs, d) => if (c == d) AONE(bs) else AZERO
  case AALTS(bs, rs) => AALTS(bs, rs.map(bder(c, _)))
  case ASEQ(bs, r1, r2) => 
    if (bnullable(r1)) AALT(bs, ASEQ(Nil, bder(c, r1), r2), fuse(bmkeps(r1), bder(c, r2)))
    else ASEQ(bs, bder(c, r1), r2)
  case ASTAR(bs, r) => ASEQ(bs, fuse(List(Z), bder(c, r)), ASTAR(Nil, r))
}

// derivative w.r.t. a string (iterates bder)
@tailrec
def bders (s: List[Char], r: ARexp) : ARexp = s match {
  case Nil => r
  case c::s => bders(s, bder(c, r))
}

// main unsimplified lexing function (produces a bitsequence)
def blex(r: ARexp, s: List[Char]) : Bits = s match {
  case Nil => if (bnullable(r)) bmkeps(r) else throw new Exception("Not matched")
  case c::cs => blex(bder(c, r), cs)
}

// calls blex and decodes the value
def blexing(r: Rexp, s: String) : Val = 
  vdecode(r, blex(internalise(r), s.toList))


// example by Tudor
//val reg = (STAR("a") ~ ("b" | "c")).%
//println(blexing(reg, "aab"))


//=======================
// simplification 
//

def flts(rs: List[ARexp]) : List[ARexp] = rs match {
  case Nil => Nil
  case AZERO :: rs => flts(rs)
  case AALTS(bs, rs1) :: rs => rs1.map(fuse(bs, _)) ++ flts(rs)
  case r1 :: rs => r1 :: flts(rs)
}

def bsimp(r: ARexp): ARexp = r match {
  case ASEQ(bs1, r1, r2) => (bsimp(r1), bsimp(r2)) match {
      case (AZERO, _) => AZERO
      case (_, AZERO) => AZERO
      case (AONE(bs2), r2s) => fuse(bs1 ++ bs2, r2s)
      // needed in order to keep the size down
      case (AALTS(bs, rs), r2) => AALTS(bs1 ++ bs, rs.map(ASEQ(Nil, _, r2)))
      case (r1s, r2s) => ASEQ(bs1, r1s, r2s)
  }
  // distinctBy deletes copies of the same  "erased" regex
  case AALTS(bs1, rs) => (flts(rs.map(bsimp))).distinctBy(erase) match {
      case Nil => AZERO
      case r::Nil => fuse(bs1, r)
      case rs => AALTS(bs1, rs)
  }
  case r => r
}

def blex_simp(r: ARexp, s: List[Char]) : Bits = s match {
  case Nil => if (bnullable(r)) bmkeps(r) 
              else throw new Exception("Not matched")
  case c::cs => blex_simp(bsimp(bder(c, r)), cs)
}

// blexing_simp decodes a value from the bitsequence (potentially slow)
// blexing2_simp decodes a string-list from the bitsequence
def blexing_simp(r: Rexp, s: String) : Val = 
  vdecode(r, blex_simp(internalise(r), s.toList))

def blexing2_simp(r: Rexp, s: String) : List[String] = 
  sdecode(r, blex_simp(internalise(r), s.toList))


//println(blexing_simp(reg, "aab"))


// extracts a string from value
def flatten(v: Val) : String = v match {
  case Empty => ""
  case Chr(c) => c.toString
  case Left(v) => flatten(v)
  case Right(v) => flatten(v)
  case Sequ(v1, v2) => flatten(v1) + flatten(v2)
  case Stars(vs) => vs.map(flatten).mkString
}

// extracts an environment from a value
def env(v: Val) : List[(String, String)] = v match {
  case Empty => Nil
  case Chr(c) => Nil
  case Left(v) => env(v)
  case Right(v) => env(v)
  case Sequ(v1, v2) => env(v1) ::: env(v2)
  case Stars(vs) => vs.flatMap(env)
  case Recd(x, v) => (x, flatten(v))::env(v)
}

def bsize(a: ARexp) = size(erase(a))

// Some Tests
//============


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

/*
val evil1 = STAR(STAR("a")) ~ "b"
val evil2 = STAR(STAR(STAR("a"))) ~ "b"
val evil3 = STAR("aa" | "a")

for(i <- 0 to 10000 by 1000) {
    println(time_needed(1, blexing2_simp(evil1, "a"*i ++ "b")))
}

for(i <- 0 to 10000 by 1000) {
    println(time_needed(1, blexing2_simp(evil2, "a"*i ++ "b")))
}

for(i <- 0 to 10000 by 1000) {
    println(time_needed(1, blexing2_simp(evil3, "a"*i)))
}
*/

// WHILE LANGUAGE
//================
def PLUS(r: Rexp) = r ~ r.%

def Range(s : List[Char]) : Rexp = s match {
  case Nil => ZERO
  case c::Nil => CHAR(c)
  case c::s => ALT(CHAR(c), Range(s))
}
def RANGE(s: String) = Range(s.toList)

val SYM = RANGE("ABCDEFGHIJKLMNOPQRSTUVXYZabcdefghijklmnopqrstuvwxyz_")
val DIGIT = RANGE("0123456789")
val ID = SYM ~ (SYM | DIGIT).% 
val NUM = PLUS(DIGIT)
val KEYWORD : Rexp = "skip" | "while" | "do" | "if" | "then" | "else" | "read" | "write" 
val SEMI: Rexp = ";"
val OP: Rexp = ":=" | "=" | "-" | "+" | "*" | "!=" | "<" | ">"
val WHITESPACE = PLUS(" " | "\n" | "\t")
val RPAREN: Rexp = "{"
val LPAREN: Rexp = "}"
val STRING: Rexp = "\"" ~ SYM.% ~ "\""


val WHILE_REGS = (("k" $ KEYWORD) | 
                  ("i" $ ID) | 
                  ("o" $ OP) | 
                  ("n" $ NUM) | 
                  ("s" $ SEMI) | 
                  ("str" $ STRING) |
                  ("p" $ (LPAREN | RPAREN)) | 
                  ("w" $ WHITESPACE)).%


// Some Simple While Tests
//========================

val prog0 = """read n"""
println(s"test: $prog0")
println(env(blexing_simp(WHILE_REGS, prog0)))
println(blexing2_simp(WHILE_REGS, prog0))

val prog1 = """read  n; write n"""  
println(s"test: $prog1")
println(env(blexing_simp(WHILE_REGS, prog1)))
println(blexing2_simp(WHILE_REGS, prog1))

val prog2 = """
write "Fib";
read n;
minus1 := 0;
minus2 := 1;
while n > 0 do {
  temp := minus2;
  minus2 := minus1 + minus2;
  minus1 := temp;
  n := n - 1
};
write "Result";
write minus2
"""

println("lexing fib program (once)")
println(blexing2_simp(WHILE_REGS, prog2).filter(s => s == "" || !s.startsWith("w")))

val n = 60
println(s"lexing fib program ($n times, size ${prog2.length * n})")
println(time_needed(1, blexing2_simp(WHILE_REGS, prog2 * n)))