app.avl.c

#include <stdio.h>
#include <stdlib.h>
#include <ntddk.h>
#include <app.avl.h>

RTL_GENERIC_COMPARE_RESULTS
CompareRoutine(
    __in struct _RTL_AVL_TABLE  *Table,
    __in PVOID  FirstStruct,
    __in PVOID  SecondStruct
)
{
    int i, j, k1, k2;
    appBLOCK A1, A2;

    A1 = ((struct appTRIPLET*)FirstStruct)->A;
    k1 = ((struct appTRIPLET*)FirstStruct)->k;
    A2 = ((struct appTRIPLET*)SecondStruct)->A;
    k2 = ((struct appTRIPLET*)SecondStruct)->k;

    if (k1 == 1 && (A1 >= A2 && A1 < A2 + k2)) return GenericEqual;
    if (k2 == 1 && (A2 >= A1 && A2 < A1 + k1)) return GenericEqual;
    if (A1 > A2) return GenericGreaterThan;
    if (A1 < A2) return GenericLessThan;
    return GenericEqual;
}

PVOID
Allocate_Routine(
    __in struct _RTL_AVL_TABLE  *Table,
    __in CLONG  ByteSize
)
{
    void *temp;
    temp = malloc(ByteSize);
    return temp;
}

VOID
FreeRoutine(
    __in struct _RTL_AVL_TABLE  *Table,
    __in PVOID  Buffer
)
{
    free(Buffer);
}

appTRIPLET*
addNode(appBLOCK A, appBLOCK B, appNUMBER k, PRTL_AVL_TABLE Table)
{
    appTRIPLET T = {A, B, k}, *ptr;
    BOOLEAN ok;
    ptr = RtlInsertElementGenericTableAvl(Table, &T, sizeof(T), &ok);
    return ptr;
}

BOOLEAN
existNode(appBLOCK A, PRTL_AVL_TABLE Table)
{
    appTRIPLET T, *P;
    T.A = A;
    T.k = 1;
    P = (appTRIPLET*)RtlLookupElementGenericTableAvl(Table, &T);
    if(P == NULL) return FALSE;
    return TRUE;
}

appTRIPLET*
findNode(appBLOCK A, PRTL_AVL_TABLE Table)
{
    appTRIPLET T, *P;
    T.A = A;
    T.k = 1;
    P = RtlLookupElementGenericTableAvl(Table, &T);
    return P;
}

BOOLEAN
deleteNode(appBLOCK A, PRTL_AVL_TABLE Table)
{
    BOOLEAN ok;
    appTRIPLET *P = findNode(A, Table);
    ok = RtlDeleteElementGenericTableAvl(Table, P);
    return ok;
}

appTRIPLET*
nextNode(appBLOCK A, PRTL_AVL_TABLE Table)
{
    appTRIPLET *ptr, *temp;
    PVOID restart;
    
    temp = findNode(A, Table);
    if(temp != NULL)
    {
        restart = (PVOID)((char*)(temp) - 0x20);
        ptr = (appTRIPLET*)RtlEnumerateGenericTableWithoutSplayingAvl(Table, &restart);
        return ptr;
    }
    temp = addNode(A, (appBLOCK)NULL, 1, Table);
    restart = (PVOID)((char*)(temp) - 0x20);
    ptr = (appTRIPLET*)RtlEnumerateGenericTableWithoutSplayingAvl(Table, &restart);
    deleteNode(A, Table);
    return ptr;
}

appBLOCK
checkNode(appBLOCK A, PRTL_AVL_TABLE Table)
{
    appTRIPLET T, *P = findNode(A, Table);
    return P->B + (A - P->A);
}

BOOLEAN
mapNode(appBLOCK A, appBLOCK B, appNUMBER k, PRTL_AVL_TABLE Table)
{
    BOOLEAN ok;
    ULONG TableLength = RtlNumberGenericTableElementsAvl(Table);
    appTRIPLET T, Left, Right, *PLeft, *PRight, *PL, *PR, *temp, *ptr;
    PVOID restart;

    //Find the edge intervals
    Left.A = A;
    Left.k = 1;
    Right.A = A + k;
    Right.k = 1;

    if(existNode(Left.A, Table))
    {
        PLeft = findNode(Left.A, Table);
        PL = PLeft;
    }
    else
    {
        PLeft = nextNode(Left.A, Table);
        PL = NULL;
    }

    if(existNode(Right.A, Table))
    {
        PR = findNode(Right.A, Table);
        PRight = nextNode(Right.A, Table);
    }
    else
    {
        PRight = nextNode(Right.A, Table);
        PR = NULL;
    }

    //Split edge intervals
    if(PL != NULL)
    {
        Left.A = PL->A;
        Left.B = PL->B;
        Left.k = A - PL->A;
        PL = &Left;
    }
    if(PR != NULL)
    {
        Right.A = A + k;
        Right.B = PR->B + ((A + k) - PR->A);
        Right.k = PR->A + PR->k - (A + k);
        PR = &Right;
    }

    // Delete crossing intervals
    restart = (PVOID)((char*)(PLeft) - 0x20);
    for(ptr = PLeft; ptr != PRight; )
    {
        temp = ptr;
        ptr = nextNode(ptr->A, Table);
        RtlDeleteElementGenericTableAvl(Table, temp);
    }

    //Add the main interval
    T.A = A;
    T.B = B;
    T.k = k;
    RtlInsertElementGenericTableAvl(Table, &T, sizeof(T), &ok);

    //Add the splitted intervals
    if(PL != NULL) addNode(PL->A, PL->B, PL->k, Table);
    if(PR != NULL) addNode(PR->A, PR->B, PR->k, Table);

    return ok;
}

BOOLEAN
unmapNode(appBLOCK A, appNUMBER k, PRTL_AVL_TABLE Table)
{
    BOOLEAN ok;
    mapNode(A, (appTRIPLET)NULL, k, Table);
    ok = deleteNode(A, Table);
    return ok;
}