#lang plait
(require (typed-in racket/base
                   [char->integer : (Char -> Number)]
                   [integer->char : (Number -> Char)]
                   [number->string : (Number -> String)]
                   [vector->list : ((Vectorof Number) -> (Listof Number))])
         (typed-in racket/vector
                   [vector-copy : ((Vectorof Number) Number Number -> (Vectorof Number))]))

(define NUMBER-TAG 1337)
(define STRING-TAG 5712)

(define MEMORY (make-vector 100 -1))
(define next-addr 0)
(define (reset) (set! next-addr 0))
(define (write-and-bump v)
  (let ([n next-addr])
    (begin
      (vector-set! MEMORY n v)
      (set! next-addr (add1 next-addr))
      n)))

(define (store-num n)
  (let ([a0 (write-and-bump NUMBER-TAG)])
    (begin
      (write-and-bump n)
      a0)))

(define (read-num a)
  (if (= (vector-ref MEMORY a) NUMBER-TAG)
      (vector-ref MEMORY (add1 a))
      (error 'number (number->string a))))

(define (store-str s)
  (let ([a0 (write-and-bump STRING-TAG)])
    (begin
      (write-and-bump (string-length s))
      (map write-and-bump
           (map char->integer (string->list s)))
      a0)))

(define (read-str a)
  (if (= (vector-ref MEMORY a) STRING-TAG)
      (let* ([len (vector-ref MEMORY (+ a 1))]
             [start (+ a 2)]
             [end (+ start len)]
             [slice (vector-copy MEMORY start end)]
             [lst (vector->list slice)])
        (list->string (map integer->char lst)))
      (error 'string (number->string a))))

(define-type-alias Value Number)
(define numV store-num)
(define strV store-str)

(define-type Exp
  [num (n : Number)]
  [str (s : String)]
  [plus (l : Exp) (r : Exp)]
  [cat (l : Exp) (r : Exp)])

(define (num+ la ra)
  (store-num (+ (read-num la) (read-num ra))))
(define (str++ la ra)
  (store-str (string-append (read-str la)
                            (read-str ra))))

(calc : (Exp -> Value))
(define (calc e)
  (type-case Exp e
    [(num n) (numV n)]
    [(str s) (strV s)]
    [(plus l r) (num+ (calc l) (calc r))]
    [(cat l r) (str++ (calc l) (calc r))]))

(define (generic-read a)
  (let ([tag (vector-ref MEMORY a)])
    (cond
      [(= tag NUMBER-TAG) (read-num a)]
      [(= tag STRING-TAG) (read-str a)]
      [else (error 'generic-read "invalid tag")])))

(test (read-num (calc (plus (num 1) (num 2)))) 3)
(test (read-num
       (calc (plus (num 1) (plus (num 2) (num 3))))) 6)
(test (read-str
       (calc (cat (str "hel") (str "lo")))) "hello")
(test (read-str
  (calc (cat (cat (str "hel")
                  (str "l")) (str "o")))) "hello")
(test/exn (calc (cat (num 1) (str "hello"))) "")
(test/exn (calc (plus (num 1) (str "hello"))) "")