#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 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-type Exp
  [num (n : Number)]
  [str (s : String)]
  [plus (l : Exp) (r : Exp)]
  [cat (l : Exp) (r : Exp)])


(define-type Type
  [stringT]
  [numT])

(define (expect2 type a b)
  (if (and (equal? type (typecheck a))
           (equal? type (typecheck b))) type
      (error 'typecheck
             (string-append "expected 2 x "
                            (to-string type)))))

(define (typecheck exp)
  (type-case Exp exp
    [(num n) (numT)]
    [(str s) (stringT)]
    [(cat a b) (expect2 (stringT) a b)]
    [(plus a b) (expect2 (numT) a b)]))

(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 (store-num n)
  (write-and-bump n))
(define (read-num a)
  (vector-ref MEMORY a))

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

(define (read-str a)
  (let* ([len (vector-ref MEMORY a)]
         [start (+ a 1)]
         [end (+ start len)]
         [slice (vector-copy MEMORY start end)]
         [lst (vector->list slice)])
    (list->string (map integer->char lst))))

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

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

(define-type result-type
  [numR (n : Number)]
  [strR (s : String)])

(define (run exp)
  (let* ([type (typecheck exp)]
         [loc (calc exp)])
    (type-case Type type
      [(numT) (numR (read-num loc))]
      [(stringT) (strR (read-str loc))])))

(test (typecheck (plus (num 1) (num 2))) (numT))
(test (run (plus (num 1) (plus (num 2) (num 3))))  (numR 6))
(test (run (cat (str "hel") (str "lo"))) (strR "hello"))
(test (typecheck
       (cat (cat (str "hel")
                 (str "l")) (str "o"))) (stringT))
(test/exn (run (cat (num 1) (str "hello"))) "expected")
(test/exn (run (plus (num 1) (str "hello"))) "expected")