// A simple lexer inspired by work of Sulzmann & Lu

//==================================================

//

// Call the test cases with 

//

//   amm lexer.sc small

//   amm lexer.sc fib

//   amm lexer.sc loops

//   amm lexer.sc email

//

//   amm lexer.sc all

// regular expressions including records

abstract class Rexp 

case object ZERO extends Rexp

case object ONE extends Rexp

case object ANYCHAR extends Rexp

case class CHAR(c: Char) extends Rexp

case class ALTS(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  

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

case class OPTIONAL(r: Rexp) extends Rexp

case class NOT(r: Rexp) extends Rexp

                // records for extracting strings or tokens

// values  

abstract class Val

case object Failure extends 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 Rec(x: String, v: Val) extends Val

case class Ntime(vs: List[Val]) extends Val

case class Optionall(v: Val) extends Val

case class Nots(s: String) extends Val

abstract class Bit

case object Z extends Bit

case object S extends Bit

type Bits = List[Bit]

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 

case class ANOT(bs: Bits, r: ARexp) extends ARexp

case class AANYCHAR(bs: Bits) extends ARexp

import scala.util.Try

trait Generator[+T] {

    self => // an alias for "this"

    def generate(): T

    def gen(n: Int) : List[T] = 

      if (n == 0) Nil else self.generate() :: gen(n - 1)

    def map[S](f: T => S): Generator[S] = new Generator[S] {

      def generate = f(self.generate())  

    }

    def flatMap[S](f: T => Generator[S]): Generator[S] = new Generator[S] {

      def generate = f(self.generate()).generate()

    }

}

  // tests a property according to a given random generator

  def test[T](r: Generator[T], amount: Int = 100)(pred: T => Boolean) {

    for (_ <- 0 until amount) {

      val value = r.generate()

      assert(pred(value), s"Test failed for: $value")

    }

    println(s"Test passed $amount times")

  }

  def test2[T, S](r: Generator[T], s: Generator[S], amount: Int = 100)(pred: (T, S) => Boolean) {

    for (_ <- 0 until amount) {

      val valueR = r.generate()

      val valueS = s.generate()

      assert(pred(valueR, valueS), s"Test failed for: $valueR, $valueS")

    }

    println(s"Test passed $amount times")

  }

// random integers

val integers = new Generator[Int] {

  val rand = new java.util.Random

  def generate() = rand.nextInt()

}

// random booleans

val booleans = integers.map(_ > 0)

// random integers in the range lo and high  

def range(lo: Int, hi: Int): Generator[Int] = 

  for (x <- integers) yield (lo + x.abs % (hi - lo)).abs

// random characters

def chars_range(lo: Char, hi: Char) = range(lo, hi).map(_.toChar)  

val chars = chars_range('a', 'z')

def oneOf[T](xs: T*): Generator[T] = 

  for (idx <- range(0, xs.length)) yield xs(idx)

def single[T](x: T) = new Generator[T] {

  def generate() = x

}   

def pairs[T, U](t: Generator[T], u: Generator[U]): Generator[(T, U)] = 

  for (x <- t; y <- u) yield (x, y)

// lists

def emptyLists = single(Nil) 

def nonEmptyLists : Generator[List[Int]] = 

  for (head <- integers; tail <- lists) yield head :: tail

def lists: Generator[List[Int]] = for {

  kind <- booleans

  list <- if (kind) emptyLists else nonEmptyLists

} yield list

def char_list(len: Int): Generator[List[Char]] = {

  if(len <= 0) single(Nil)

  else{

    for { 

      c <- chars_range('a', 'c')

      tail <- char_list(len - 1)

    } yield c :: tail

  }

}

def strings(len: Int): Generator[String] = for(cs <- char_list(len)) yield cs.toString

def sampleString(r: Rexp) : List[String] = r match {

  case STAR(r) => stringsFromRexp(r).flatMap(s => List("", s, s ++ s))//only generate 0, 1, 2 reptitions

  case SEQ(r1, r2) => stringsFromRexp(r1).flatMap(s1 => stringsFromRexp(r2).map(s2 => s1 ++ s2) )

  case ALTS(r1, r2) => throw new Error(s" Rexp ${r} not expected: all alternatives are supposed to have been opened up")

  case ONE => "" :: Nil

  case ZERO => Nil

  case CHAR(c) => c.toString :: Nil

}

def stringsFromRexp(r: Rexp) : List[String] = 

  breakIntoTerms(r).flatMap(r => sampleString(r))

// (simple) binary trees

trait Tree[T]

case class Inner[T](left: Tree[T], right: Tree[T]) extends Tree[T]

case class Leaf[T](x: T) extends Tree[T]

def leafs[T](t: Generator[T]): Generator[Leaf[T]] = 

  for (x <- t) yield Leaf(x)

def inners[T](t: Generator[T]): Generator[Inner[T]] = 

  for (l <- trees(t); r <- trees(t)) yield Inner(l, r)

def trees[T](t: Generator[T]): Generator[Tree[T]] = 

  for (kind <- range(0, 2);  

       tree <- if (kind == 0) leafs(t) else inners(t)) yield tree

// regular expressions

// generates random leaf-regexes; prefers CHAR-regexes

def leaf_rexp() : Generator[Rexp] =

  for (kind <- range(0, 5);

       c <- chars_range('a', 'd')) yield

    kind match {

      case 0 => ZERO

      case 1 => ONE

      case _ => CHAR(c) 

    }

// generates random inner regexes with maximum depth d

def inner_rexp(d: Int) : Generator[Rexp] =

  for (kind <- range(0, 3);

       l <- rexp(d); 

       r <- rexp(d))

  yield kind match {

    case 0 => ALTS(l, r)

    case 1 => SEQ(l, r)

    case 2 => STAR(r)

  }

// generates random regexes with maximum depth d;

// prefers inner regexes in 2/3 of the cases

def rexp(d: Int = 100): Generator[Rexp] = 

  for (kind <- range(0, 3);

       r <- if (d <= 0 || kind == 0) leaf_rexp() else inner_rexp(d - 1)) yield r

// some test functions for rexps

def height(r: Rexp) : Int = r match {

  case ZERO => 1

  case ONE => 1

  case CHAR(_) => 1

  case ALTS(r1, r2) => 1 + List(height(r1), height(r2)).max

  case SEQ(r1, r2) =>  1 + List(height(r1), height(r2)).max

  case STAR(r) => 1 + height(r)

}

// def size(r: Rexp) : Int = r match {

//   case ZERO => 1

//   case ONE => 1

//   case CHAR(_) => 1

//   case ALTS(r1, r2) => 1 + size(r1) + size(r2)

//   case SEQ(r1, r2) =>  1 + size(r1) + size(r2)

//   case STAR(r) => 1 + size(r) 

// }

// randomly subtracts 1 or 2 from the STAR case

def size_faulty(r: Rexp) : Int = r match {

  case ZERO => 1

  case ONE => 1

  case CHAR(_) => 1

  case ALTS(r1, r2) => 1 + size_faulty(r1) + size_faulty(r2)

  case SEQ(r1, r2) =>  1 + size_faulty(r1) + size_faulty(r2)

  case STAR(r) => 1 + size_faulty(r) - range(0, 2).generate

}

// 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) = ALTS(r, s)

  def % = STAR(r)

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

}

implicit def stringOps(s: String) = new {

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

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

  def % = STAR(s)

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

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

  def $ (r: Rexp) = RECD(s, r)

}

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

  case ZERO => false

  case ONE => true

  case CHAR(_) => false

  case ANYCHAR => false

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

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

  case STAR(_) => true

  case RECD(_, r1) => nullable(r1)

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

  case OPTIONAL(r) => true

  case NOT(r) => !nullable(r)

}

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 ANYCHAR => ONE 

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

  case SEQ(r1, r2) => 

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

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

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

  case RECD(_, r1) => der(c, r1)

  case NTIMES(n, r) => if(n > 0) SEQ(der(c, r), NTIMES(n - 1, r)) else ZERO

  case OPTIONAL(r) => der(c, r)

  case NOT(r) =>  NOT(der(c, r))

}

// extracts a string from a 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

  case Ntime(vs) => vs.map(flatten).mkString

  case Optionall(v) => flatten(v)

  case Rec(_, v) => flatten(v)

}

// extracts an environment from a value;

// used for tokenising a string

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 Ntime(vs) => vs.flatMap(env)

  case Rec(x, v) => (x, flatten(v))::env(v)

  case Optionall(v) => env(v)

  case Nots(s) => ("Negative", s) :: Nil

}

// The injection and mkeps part of the lexer

//===========================================

def mkeps(r: Rexp) : Val = r match {

  case ONE => Empty

  case ALTS(r1, r2) => 

    if (nullable(r1)) Left(mkeps(r1)) else Right(mkeps(r2))

  case SEQ(r1, r2) => Sequ(mkeps(r1), mkeps(r2))

  case STAR(r) => Stars(Nil)

  case RECD(x, r) => Rec(x, mkeps(r))

  case NTIMES(n, r) => Ntime(List.fill(n)(mkeps(r)))

  case OPTIONAL(r) => Optionall(Empty)

  case NOT(rInner) => if(nullable(rInner)) throw new Exception("error")  

                         else Nots("")//Nots(s.reverse.toString)

//   case NOT(ZERO) => Empty

//   case NOT(CHAR(c)) => Empty

//   case NOT(SEQ(r1, r2)) => Sequ(mkeps(NOT(r1)), mkeps(NOT(r2)))

//   case NOT(ALTS(r1, r2)) => if(!nullable(r1)) Left(mkeps(NOT(r1))) else Right(mkeps(NOT(r2)))

//   case NOT(STAR(r)) => Stars(Nil) 

}

def inj(r: Rexp, c: Char, v: Val) : Val = (r, v) match {

  case (STAR(r), Sequ(v1, Stars(vs))) => Stars(inj(r, c, v1)::vs)

  case (SEQ(r1, r2), Sequ(v1, v2)) => Sequ(inj(r1, c, v1), v2)

  case (SEQ(r1, r2), Left(Sequ(v1, v2))) => Sequ(inj(r1, c, v1), v2)

  case (SEQ(r1, r2), Right(v2)) => Sequ(mkeps(r1), inj(r2, c, v2))

  case (ALTS(r1, r2), Left(v1)) => Left(inj(r1, c, v1))

  case (ALTS(r1, r2), Right(v2)) => Right(inj(r2, c, v2))

  case (CHAR(d), Empty) => Chr(c) 

  case (RECD(x, r1), _) => Rec(x, inj(r1, c, v))

  case (NTIMES(n, r), Sequ(v1, Ntime(vs))) => Ntime(inj(r, c, v1)::vs)

  case (OPTIONAL(r), v) => Optionall(inj(r, c, v))

  case (NOT(r), Nots(s)) => Nots(c.toString ++ s)

  case (ANYCHAR, Empty) => Chr(c)

}

    def size(r: Rexp) : Int = r match {

      case ZERO => 1

      case ONE => 1

      case CHAR(_) => 1

      case ANYCHAR => 1

      case NOT(r0) => 1 + size(r0)

      case SEQ(r1, r2) => 1 + size(r1) + size(r2)

      case ALTS(r1, r2) => 1 + List(r1, r2).map(size).sum

      case STAR(r) => 1 + size(r)

    }

//almost as if a zeroable in called on each subexpression, and if so, turn that sub-expression into 0

def light_simp(r: Rexp): Rexp = r match {

  case ALTS(r1, r2) => (light_simp(r1), light_simp(r2)) match {

    case (ZERO, ZERO) => ZERO

    case (r1p, r2p) => ALTS(r1p, r2p)

  }

  case SEQ(r1, r2) => (light_simp(r1), light_simp(r2)) match {

    case (ZERO, _) => ZERO

    case (_, ZERO) => ZERO

    case (r1p, r2p) => SEQ(r1p, r2p)

  }

  case STAR(r) => STAR(r)

  case CHAR(d) => CHAR(d)

  case ONE => ONE

  case ZERO => ZERO

}

def simp_lexer(r: Rexp, s: List[Char]): Val = s match{

  case Nil => {//println("mkeps!"); println(r); 

    val v = mkeps(r); 

    //println(v); 

    v}

  case c::cs => { 

    val derc1 = der(c, r);

    val derc = light_simp(derc1);

    println("derivative size: "+ size(derc));

    val v = simp_lexer(derc, cs); 

    //println("Injection!");println("r: "+r); println("c: "+c); println("v:"+v); 

    inj(r, c, v ) }

}

def print_lexer(r: Rexp, s: List[Char]): Val = s match{

  case Nil => {//println("mkeps!"); println(r); 

    val v = mkeps(r); 

    //println(v); 

    v}

  case c::cs => { 

    val derc = der(c, r);

    //val derc = light_simp(derc1);

    println("derivative size: "+ size(derc));

    val v = print_lexer(derc, cs); 

    //println("Injection!");println("r: "+r); println("c: "+c); println("v:"+v); 

    inj(r, c, v ) }

}

val r = STAR("a"|"aa") ~ ("c")

println("unsimplified lexer's intermediate derivative sizes:")

val v = print_lexer(r, "aac".toList)

println("simplified lexer's intermediate derivative sizes:")

val v2 = simp_lexer(r, "aac".toList)

//println(v)

// some "rectification" functions for simplification

// The Lexing Rules for the WHILE Language

  // bnullable function: tests whether the aregular 

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

    case ANOT(_, rn) => !bnullable(rn)

  }

def mkepsBC(r: ARexp) : Bits = r match {

    case AONE(bs) => bs

    case AALTS(bs, rs) => {

      val n = rs.indexWhere(bnullable)

      bs ++ mkepsBC(rs(n))

    }

    case ASEQ(bs, r1, r2) => bs ++ mkepsBC(r1) ++ mkepsBC(r2)

    case ASTAR(bs, r) => bs ++ List(Z)

    case ANOT(bs, rn) => bs

  }

def bder(c: Char, r: ARexp) : ARexp = r match {

    case AZERO => AZERO

    case AONE(_) => AZERO

    case ACHAR(bs, f) => if (c == f) AONE(bs) else AZERO

    case AALTS(bs, rs) => AALTS(bs, rs.map(bder(c, _)))

    case ASEQ(bs, r1, r2) => 

      if (bnullable(r1)) AALTS(bs, ASEQ(Nil, bder(c, r1), r2) :: fuse(mkepsBC(r1), bder(c, r2)) :: Nil )

      else ASEQ(bs, bder(c, r1), r2)

    case ASTAR(bs, r) => ASEQ(bs, fuse(List(S), bder(c, r)), ASTAR(Nil, r))

    case ANOT(bs, rn) => ANOT(bs, bder(c, rn))

    case AANYCHAR(bs) => AONE(bs)

  } 

def fuse(bs: Bits, r: ARexp) : ARexp = r match {

    case AZERO => AZERO

    case AONE(cs) => AONE(bs ++ cs)

    case ACHAR(cs, f) => ACHAR(bs ++ cs, f)

    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)

    case ANOT(cs, r) => ANOT(bs ++ cs, r)

  }

def internalise(r: Rexp) : ARexp = r match {

    case ZERO => AZERO

    case ONE => AONE(Nil)

    case CHAR(c) => ACHAR(Nil, c)

    //case PRED(f) => APRED(Nil, f)

    case ALTS(r1, r2) => 

      AALTS(Nil, List(fuse(List(Z), internalise(r1)), fuse(List(S), internalise(r2))))

    // case ALTS(r1::rs) => {

    //   val AALTS(Nil, rs2) = internalise(ALTS(rs))

    //   AALTS(Nil, fuse(List(Z), internalise(r1)) :: rs2.map(fuse(List(S), _)))

    // }

    case SEQ(r1, r2) => ASEQ(Nil, internalise(r1), internalise(r2))

    case STAR(r) => ASTAR(Nil, internalise(r))

    case RECD(x, r) => internalise(r)

    case NOT(r) => ANOT(Nil, internalise(r))

    case ANYCHAR => AANYCHAR(Nil)

  }

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)

          case (r1s, r2s) => ASEQ(bs1, r1s, r2s)

      }

      case AALTS(bs1, rs) => {

            val rs_simp = rs.map(bsimp(_))

            val flat_res = flats(rs_simp)

            val dist_res = distinctBy(flat_res, erase)//strongDB(flat_res)//distinctBy(flat_res, erase)

            dist_res match {

              case Nil => AZERO

              case s :: Nil => fuse(bs1, s)

              case rs => AALTS(bs1, rs)  

            }

      }

      case r => r

    }

}

def strongBsimp(r: ARexp): ARexp =

{

  r match {

    case ASEQ(bs1, r1, r2) => (strongBsimp(r1), strongBsimp(r2)) match {

        case (AZERO, _) => AZERO