#!/usr/bin/env nomsu -V2.5.4.3
#
    This file contains code that supports manipulating and using collections like lists
    and dictionaries.

use "core/metaprogramming.nom"
use "core/control_flow.nom"
use "core/operators.nom"

# List/dict functions:
# Indexing
test:
    assume ((2 nd to last in [1, 2, 3, 4, 5]) is 4)

compile [..]
    %index st to last in %list, %index nd to last in %list
    %index rd to last in %list, %index th to last in %list
..to (Lua value "utils.nth_to_last(\(%list as lua expr), \(%index as lua expr))")

~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~

parse [last in %list] as (1 st to last in %list)
parse [first in %list] as %list.1

# Membership testing
test:
    assume (3 is in [1, 2, 3, 4, 5])

action [%item is in %list, %list contains %item, %list has %item]:
    for %key = %value in %list:
        if (%key is %item): return (yes)
    return (no)

test:
    assume (99 isn't in [1, 2, 3])

action [..]
    %item isn't in %list, %item is not in %list, %list doesn't contain %item
    %list does not contain %item, %list doesn't have %item, %list does not have %item
..:
    for %key = %value in %list:
        if (%key is %item): return (no)
    return (yes)

test:
    assume ({x:no} has key "x")

parse [%list has key %index, %list has index %index] as (%list.%index != (nil))
test:
    assume ({x:no} doesn't have key "y")
    assume (not ({x:no} doesn't have key "x"))

parse [..]
    %list doesn't have key %index, %list does not have key %index
    %list doesn't have index %index, %list does not have index %index
..as (%list.%index == (nil))

compile [number of keys in %list] to (..)
    Lua value "utils.size(\(%list as lua expr))"

test:
    %list = [1, 2, 3, 4, 5]
    append 6 to %list
    assume ((last in %list) is 6)
    pop from %list
    assume ((last in %list) is 5)
    remove index 1 from %list
    assume ((first in %list) is 2)

compile [..]
    append %item to %list, add %item to %list, to %list add %item, to %list append %item
..to (Lua "table.insert(\(%list as lua expr), \(%item as lua expr))")

compile [add %item to %list at index %i] to (..)
    Lua ".."
        table.insert(\(%list as lua expr), \(%i as lua expr), \(%item as lua expr))

compile [pop from %list, remove last from %list] to (..)
    Lua value "table.remove(\(%list as lua expr))"

compile [remove index %index from %list] to (..)
    Lua value "table.remove(\(%list as lua expr), \(%index as lua expr))"

~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~

# List Comprehension
test:
    assume (((% * %) for % in [1, 2, 3]) == [1, 4, 9])

parse [%expression for %item in %iterable] as (..)
    result of:
        %comprehension = []
        for %item in %iterable:
            add %expression to %comprehension
        return %comprehension

parse [..]
    %expression for %index in %start to %stop via %step
    %expression for %index in %start to %stop by %step
..as (..)
    result of:
        %comprehension = []
        for %index in %start to %stop via %step:
            add %expression to %comprehension
        return %comprehension

~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~

test:
    assume (((% * %) for % in 1 to 3) == [1, 4, 9])

parse [%expression for %var in %start to %stop] as (..)
    %expression for %var in %start to %stop via 1

test:
    assume (("\%k,\%v" for %k = %v in {x:1}) == ["x,1"])

parse [..]
    %expression for %key = %value in %iterable
    %expression for %key %value in %iterable
..as (..)
    result of:
        %comprehension = []
        for %key = %value in %iterable:
            add %expression to %comprehension
        return %comprehension

# Dict comprehensions
test:
    assume (((% * %) = % for % in [1, 2, 3]) == {1:1, 4:2, 9:3})

parse [%key = %value for %item in %iterable, %key %value for %item in %iterable] as
..(..)
    result of:
        %comprehension = {}
        for %item in %iterable:
            %comprehension.%key = %value
        return %comprehension

test:
    assume ((%k = (%v * %v) for %k = %v in {x:1, y:2, z:3}) == {x:1, y:4, z:9})

parse [..]
    %key = %value for %src_key = %src_value in %iterable
    %key %value for %src_key %src_value in %iterable
..as (..)
    result of:
        %comprehension = {}
        for %src_key = %src_value in %iterable:
            %comprehension.%key = %value
        return %comprehension

parse [..]
    %key = %value for %item in %start to %stop via %step
    %key %value for %item in %start to %stop via %step
..as (..)
    result of:
        %comprehension = {}
        for %item in %start to %stop via %step:
            %comprehension.%key = %value
        return %comprehension

~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~

test:
    assume (((% * %) = % for % in 1 to 3) == {1:1, 4:2, 9:3})

parse [..]
    %key = %value for %item in %start to %stop
    %key %value for %item in %start to %stop
..as (%key = %value for %item in %start to %stop via 1)

test:
    assume (([[1, 2], [3, 4]] flattened) == [1, 2, 3, 4])

action [%lists flattened]:
    %flat = []
    for %list in %lists:
        for %item in %list: add %item to %flat
    return %flat

test:
    assume ((entries in {x:1}) == [{key:"x", value:1}])

parse [entries in %dict] as ({key:%k, value:%v} for %k = %v in %dict)
test:
    assume ((keys in {x:1}) == ["x"])

parse [keys in %dict, keys of %dict] as (%k for %k = %v in %dict)
test:
    assume ((values in {x:1}) == [1])

parse [values in %dict, values of %dict] as (%v for %k = %v in %dict)

# Metatable stuff
test:
    %t = {}
    set %t 's metatable to {__tostring:[%] -> "XXX"}
    assume ("\%t" == "XXX")

compile [set %dict 's metatable to %metatable] to (..)
    Lua "setmetatable(\(%dict as lua expr), \(%metatable as lua expr));"

test:
    assume (({} with fallback % -> (% + 1)).10 == 11)

compile [%dict with fallback %key -> %value] to (..)
    Lua value ".."
        setmetatable(\(%dict as lua expr), {__index=function(self, \(%key as lua expr))
            local value = \(%value as lua expr)
            self[\(%key as lua expr)] = value
            return value
        end})

# Sorting
test:
    %x = [3, 1, 2]
    sort %x
    assume (%x == [1, 2, 3])
    sort %x by % = (- %)
    assume (%x == [3, 2, 1])
    %keys = {1:999, 2:0, 3:50}
    sort %x by % = %keys.%
    assume (%x == [2, 3, 1])

compile [sort %items] to (Lua "table.sort(\(%items as lua expr));")
parse [sort %items by %item = %key_expr, sort %items by %item -> %key_expr] as (..)
    do:
        %keys = ({} with fallback %item -> %key_expr)
        lua> "table.sort(\%items, function(x,y) return \%keys[x] < \%keys[y] end)"

~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~

test:
    assume ((sorted [3, 1, 2]) == [1, 2, 3])

action [%items sorted, sorted %items]:
    %copy = (% for % in %items)
    sort %copy
    return %copy

parse [%items sorted by %item = %key, %items sorted by %item -> %key] as (..)
    result of:
        %copy = (% for % in %items)
        sort %copy by %item = %key
        return %copy

test:
    assume ((unique [1, 2, 1, 3, 2, 3]) == [1, 2, 3])

action [unique %items]:
    %unique = []
    %seen = {}
    for % in %items:
        unless %seen.%:
            add % to %unique
            %seen.% = (yes)
    
    return %unique