#lang typed/racket #:optimize
(require (prefix-in sh: scheme/base))
(provide filter remove List reverse
empty empty? list-length cons head tail
(rename-out [first* first] [rest* rest] [list-map map]
[list-foldr foldr] [list-foldl foldl])
list-ref list-set drop ->list list)
(struct: (A) Leaf ([first : A]))
(struct: (A) Node ([first : A]
[left : (Tree A)]
[right : (Tree A)]))
(define-type (Tree A) (U (Leaf A) (Node A)))
(struct: (A) Root ([size : Integer]
[first : (Tree A)]
[rst : (RAList A)]))
(define-type (RAList A) (U Null (Root A)))
(define-type (List A) (RAList A))
(define empty-tree (Leaf null))
(define empty null)
(: empty? : (All (A) ((RAList A) -> Boolean)))
(define (empty? list)
(null? list))
(: list-length : (All (A) ((RAList A) -> Integer)))
(define (list-length list)
(if (null? list) 0 (+ (Root-size list) (list-length (Root-rst list)))))
(: cons : (All (A) (A (RAList A) -> (RAList A))))
(define (cons elem list)
(if (null? list)
(Root 1 (Leaf elem) list)
(let* ([rst (Root-rst list)]
[lsize (list-length list)]
[rst-size (list-length rst)])
(if (eq? lsize rst-size)
(Root (+ 1 lsize rst-size)
(Node elem (Root-first list) (first rst))
(rest rst))
(Root 1 (Leaf elem) list)))))
(: first : (All (A) ((RAList A) -> (Tree A))))
(define (first list )
(if (null? list)
(error 'cons "given list is empty") (Root-first list)))
(: rest : (All (A) ((RAList A) -> (RAList A))))
(define (rest list)
(if (null? list)
(error 'cons "given list is empty") (Root-rst list)))
(: head : (All (A) ((RAList A) -> A)))
(define (head list)
(if (null? list)
(error 'head "given list is empty")
(let ([first (Root-first list)])
(if (Leaf? first)
(Leaf-first first)
(Node-first first)))))
(: tail : (All (A) ((RAList A) -> (RAList A))))
(define (tail list)
(if (null? list)
(error 'tail "given list is empty")
(let ([first (Root-first list)]
[rst (Root-rst list)]
[size (arithmetic-shift (Root-size list) -1)])
(if (Leaf? first)
rst
(Root size (Node-left first) (Root size (Node-right first) rst))))))
(: tree-lookup : (All (A) (Integer (Tree A) Integer -> A)))
(define (tree-lookup size tre pos)
(let ([pos-zero? (zero? pos)])
(cond
[(and (Leaf? tre) pos-zero?) (Leaf-first tre)]
[(Node? tre) (tree-lookup-help size tre pos pos-zero?)]
[else (error 'list-ref "given index out of bound")])))
(: tree-lookup-help : (All (A) (Integer (Node A) Integer Boolean -> A)))
(define (tree-lookup-help size tre pos pos-zero?)
(let ([newsize (arithmetic-shift size -1)])
(cond
[pos-zero? (Node-first tre)]
[(<= pos newsize) (tree-lookup newsize (Node-left tre) (- pos 1))]
[else (tree-lookup newsize (Node-right tre) (- pos 1 newsize))])))
(: tree-update : (All (A) (Integer (Tree A) Integer A -> (Tree A))))
(define (tree-update size tre pos elem)
(let ([newsize (arithmetic-shift size -1)]
[pos-zero? (zero? pos)])
(cond
[(and (Leaf? tre) pos-zero?) (Leaf elem)]
[(Node? tre) (tree-update-helper newsize tre pos elem pos-zero?)]
[else (error 'list-set "given index out of bound")])))
(: tree-update-helper :
(All (A) (Integer (Node A) Integer A Boolean -> (Tree A))))
(define (tree-update-helper newsize tre pos elem pos-zero?)
(let ([left (Node-left tre)]
[right (Node-right tre)]
[first (Node-first tre)])
(cond
[(eq? pos 0) (Node elem left right)]
[(<= pos newsize)
(Node first (tree-update newsize left (- pos 1) elem) right)]
[else (Node first left
(tree-update newsize right (- pos 1 newsize) elem))])))
(: list-ref : (All (A) ((RAList A) Integer -> A)))
(define (list-ref list pos)
(if (null? list)
(error 'list-ref "given index out of bound")
(let ([size (Root-size list)])
(if (< pos size)
(tree-lookup size (Root-first list) pos)
(list-ref (Root-rst list) (- pos size))))))
(: list-set : (All (A) ((RAList A) Integer A -> (RAList A))))
(define (list-set list pos elem)
(if (null? list)
(error 'list-set "given index out of bound")
(let ([size (Root-size list)]
[first (Root-first list)]
[rst (Root-rst list)])
(if (< pos size)
(Root size (tree-update size first pos elem) rst)
(Root size first (list-set rst (- pos size) elem))))))
(: tree-drop : (All (A) (Integer (Tree A) Integer (RAList A) -> (RAList A))))
(define (tree-drop size tre pos list)
(cond
[(zero? pos) (Root size tre list)]
[(and (Leaf? tre) (= pos 1)) list]
[(Node? tre) (tree-drop-help size tre pos list)]
[else (error 'drop "not enough elements to drop")]))
(: tree-drop-help :
(All (A) (Integer (Node A) Integer (RAList A) -> (RAList A))))
(define (tree-drop-help size tre pos list)
(let ([newsize (arithmetic-shift size -1)]
[left (Node-left tre)]
[right (Node-right tre)])
(if (<= pos newsize)
(tree-drop newsize left (sub1 pos)
(Root newsize right list))
(tree-drop newsize right (- pos 1 newsize) list))))
(: drop : (All (A) (Integer (RAList A) -> (RAList A))))
(define (drop pos list)
(cond
[(zero? pos) list]
[(Root? list) (drop-help list pos)]
[else (error 'drop "not enough elements to drop")]))
(: drop-help : (All (A) ((Root A) Integer -> (RAList A))))
(define (drop-help list pos)
(let ([size (Root-size list)]
[first (Root-first list)]
[rst (Root-rst list)])
(if (< pos size)
(tree-drop size first pos rst)
(drop (- pos size) rst))))
(: list-map :
(All (A C B ...)
(case-lambda
((A -> C) (List A) -> (List C))
((A B ... B -> C) (List A) (List B) ... B -> (List C)))))
(define list-map
(pcase-lambda: (A C B ...)
[([func : (A -> C)]
[list : (List A)])
(if (empty? list)
empty
(cons (func (head list)) (list-map func (tail list))))]
[([func : (A B ... B -> C)]
[list : (List A)] . [lists : (List B) ... B])
(if (or (empty? list) (ormap empty? lists))
empty
(cons (apply func (head list) (map head lists))
(apply list-map
func
(tail list)
(map tail lists))))]))
(: list-foldr :
(All (A C B ...)
(case-lambda ((C A -> C) C (List A) -> C)
((C A B ... B -> C) C (List A) (List B) ... B -> C))))
(define list-foldr
(pcase-lambda: (A C B ...)
[([func : (C A -> C)]
[base : C]
[list : (List A)])
(if (empty? list)
base
(func (list-foldr func base (tail list)) (head list)))]
[([func : (C A B ... B -> C)]
[base : C]
[list : (List A)] . [lists : (List B) ... B])
(if (or (empty? list) (ormap empty? lists))
base
(apply func (apply list-foldr
func
base
(tail list)
(map tail lists))
(head list)
(map head lists)))]))
(: list-foldl :
(All (A C B ...)
(case-lambda ((C A -> C) C (List A) -> C)
((C A B ... B -> C) C (List A) (List B) ... B -> C))))
(define list-foldl
(pcase-lambda: (A C B ...)
[([func : (C A -> C)]
[base : C]
[list : (List A)])
(if (empty? list)
base
(list-foldl func (func base (head list)) (tail list)))]
[([func : (C A B ... B -> C)]
[base : C]
[list : (List A)] . [lists : (List B) ... B])
(if (or (empty? list) (ormap empty? lists))
base
(apply list-foldl
func
(apply func base (head list) (map head lists))
(tail list)
(map tail lists)))]))
(: ->list : (All (A) ((RAList A) -> (Listof A))))
(define (->list list)
(if (empty? list)
null
(sh:cons (head list) (->list (tail list)))))
(: list : (All (A) (A * -> (RAList A))))
(define (list . rst)
(foldr (inst cons A) empty rst))
(define first* head)
(define rest* tail)
(: filter : (All (A) ((A -> Boolean) (List A) -> (List A))))
(define (filter func ral)
(if (empty? ral)
empty
(let ([head (head ral)]
[tail (tail ral)])
(if (func head)
(cons head (filter func tail))
(filter func tail)))))
(: remove : (All (A) ((A -> Boolean) (List A) -> (List A))))
(define (remove func ral)
(if (empty? ral)
empty
(let ([head (head ral)]
[tail (tail ral)])
(if (func head)
(remove func tail)
(cons head (remove func tail))))))
(: reverse : (All (A) ((List A) -> (List A))))
(define (reverse ral)
(: local-reverse : (All (A) ((List A) (List A) -> (List A))))
(define (local-reverse ral accum)
(if (empty? ral)
accum
(local-reverse (tail ral) (cons (head ral) accum))))
(local-reverse ral empty))