-- glslfx version 0.1

//
// Copyright 2018 Pixar
//
// Licensed under the Apache License, Version 2.0 (the "Apache License")
// with the following modification; you may not use this file except in
// compliance with the Apache License and the following modification to it:
// Section 6. Trademarks. is deleted and replaced with:
//
// 6. Trademarks. This License does not grant permission to use the trade
//    names, trademarks, service marks, or product names of the Licensor
//    and its affiliates, except as required to comply with Section 4(c) of
//    the License and to reproduce the content of the NOTICE file.
//
// You may obtain a copy of the Apache License at
//
//     http://www.apache.org/licenses/LICENSE-2.0 
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the Apache License with the above modification is
// distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
// KIND, either express or implied. See the Apache License for the specific
// language governing permissions and limitations under the Apache License.
//

--- This is what an import might look like.
--- #import $TOOLS/hdx/shaders/selection.glslfx

--- --------------------------------------------------------------------------
-- glsl Selection.DecodeUtils
/// Decoding utilities for selection highlighting.
/// This mixin may be safely included by any shader stage. Its lone dependency 
/// on code generation is the hdxSelectionBuffer SSBO that holds the encoded
/// selection state.

// This should match the HdSelection enum.
int GetMaxNumSelectionHighlightModes()
{
    return 2;
}

#if defined(HD_HAS_hdxSelectionBuffer)
// helper methods to decode hdxSelectionBuffer
// ------ selection highlight mode specifics -------
bool HighlightModeHasSelection(int offset)
{
    return (offset != 0);
}

int GetIndexOffsetForHighlightMode(int mode)
{
    return HdGet_hdxSelectionBuffer(mode + 1/*[0] holds #modes*/);
}

int GetNumSelectionHighlightModes()
{
    return HdGet_hdxSelectionBuffer(0);
}

// ------  selection offset specifics -------
// Note: This should match the selection offset encoding in
// hdx/SelectionTracker.cpp
void DecodeSelOffset(int selOffset, REF(thread, bool) isSelected, 
                     REF(thread, int) nextOffset)
{
    isSelected = bool(selOffset & 0x1); // bit 0
    nextOffset = selOffset >> 1;        // bits 31:1
}

// --------- subprim decoding ----------
// Each subprims offsets' buffer encoding is:
// [subprim-type][min][max][     selOffsets     ]
// <----------3 ----------><--- max - min + 1 -->
struct SubprimHeader {
    int type;
    int min;
    int max;
};
SubprimHeader DecodeSubprimHeader(int offset)
{
    SubprimHeader header;
    header.type = HdGet_hdxSelectionBuffer(offset    );
    header.min  = HdGet_hdxSelectionBuffer(offset + 1);
    header.max  = HdGet_hdxSelectionBuffer(offset + 2);
    return header;
}

bool IsSubprimSelected(int id, int min, int max, int headerStart,
                       REF(thread, int) nextSubprimOffset)
{
    const int SUBPRIM_SELOFFSETS_HEADER_SIZE  = 3;

    nextSubprimOffset = 0; // initialize
    bool isSelected = false;

    if (id >= min && id < max) {
        int netSubprimIndex =
            headerStart + SUBPRIM_SELOFFSETS_HEADER_SIZE + id - min;
        int selOffset = HdGet_hdxSelectionBuffer(netSubprimIndex);

        DecodeSelOffset(selOffset, /*out*/isSelected, /*out*/nextSubprimOffset);
    } else {
        // The subprim id does not fall in the selected id range, so the subprim
        // in question isn't selected. However, we can still have other
        // subprim(s) selected. To get the offset to jump to, we mimic decoding
        // the "min" subprim id.
        int minSubprimIndex = headerStart + SUBPRIM_SELOFFSETS_HEADER_SIZE;
        DecodeSelOffset(HdGet_hdxSelectionBuffer(minSubprimIndex),
                        /*out*/isSelected, /*out*/nextSubprimOffset);
        isSelected = false;
    }

    return isSelected;
}
#endif // HD_HAS_hdxSelectionBuffer

// Keep the result struct definition outside guards.
struct SubprimSelectionResult {
    bool elementSelected;
    bool edgeSelected;
    bool pointSelected;
    int  pointColorId;
};

SubprimSelectionResult InitSubprimSelectionResult(
    bool _elementSelected,
    bool _edgeSelected,
    bool _pointSelected,
    int _pointColorId)
{
    SubprimSelectionResult res;

    res.elementSelected = _elementSelected;
    res.edgeSelected = _edgeSelected;
    res.pointSelected = _pointSelected;
    res.pointColorId = _pointColorId;

    return res;
}

bool IsAnySubprimSelected(SubprimSelectionResult s)
{
    return (s.elementSelected || s.edgeSelected || s.pointSelected);
}
bool HasCustomizedPointColor(SubprimSelectionResult s)
{
    // Use -1 to encode selected points that don't have a custom point color.
    return (s.pointColorId != -1);
}

#if defined(HD_HAS_hdxSelectionBuffer)
// Note: These should match the SubprimType enum in hdx/SelectionTracker.cpp
#define SUBPRIM_TYPE_ELEMENT 0
#define SUBPRIM_TYPE_EDGE    1
#define SUBPRIM_TYPE_POINT   2
#define SUBPRIM_TYPE_INSTANCE 3

SubprimSelectionResult GetSubprimSel(int offset,
                                     int elementId, int edgeId, int pointId)
{
    SubprimSelectionResult s = InitSubprimSelectionResult(false, false, false, 0);

    int nextSubprimOffset = 0;
    SubprimHeader header = DecodeSubprimHeader(offset);

    if (header.type == SUBPRIM_TYPE_ELEMENT) {
        s.elementSelected =
            IsSubprimSelected(elementId, header.min, header.max,
                              offset, /*out*/nextSubprimOffset);

        if (nextSubprimOffset != 0) {
            // fragment has additional subprim(s) selected. update header.
            header = DecodeSubprimHeader(nextSubprimOffset);
            offset = nextSubprimOffset;
        }
    }

    if (header.type == SUBPRIM_TYPE_EDGE) {
        s.edgeSelected =
            IsSubprimSelected(edgeId, header.min, header.max,
                              offset, /*out*/nextSubprimOffset);

        if (nextSubprimOffset != 0) {
            // fragment has points selected. update header.
            header = DecodeSubprimHeader(nextSubprimOffset);
            offset = nextSubprimOffset;
        }
    }

    if (header.type == SUBPRIM_TYPE_POINT) {
        s.pointSelected =
            IsSubprimSelected(pointId, header.min, header.max,
                              offset, /*unused*/nextSubprimOffset);
        // For points alone, since there isn't any subprim to follow it, the
        // offset field is overriden to represent the index into the
        // selectedPointColors buffer to support customized coloring of a set of
        // selected points.
        s.pointColorId = nextSubprimOffset;
    }

    return s;
}

// --------- instance decoding ----------
bool IsInstanceSelected(int offset, REF(thread, int) nextOffset)
{
    // If we don't find an instance subprim block, pass the same offset to
    // GetSubprimSel.
    nextOffset = offset;
    bool sel = false;

    int instanceId = GetDrawingCoord().instanceIndex[0];

    SubprimHeader header = DecodeSubprimHeader(offset);
    if (header.type == SUBPRIM_TYPE_INSTANCE) {
        sel =
            IsSubprimSelected(instanceId, header.min, header.max,
                              offset, /*out*/nextOffset);
    }

    return sel;
}
#endif // HD_HAS_hdxSelectionBuffer

// --------- selection buffer decoding entry point ----------
struct SelectionResult {
    bool primOrInstanceSel;
    SubprimSelectionResult subprimSel;
};

SelectionResult InitSelectionResult(
    bool _primOrInstanceSel,
    bool _elementSelected,
    bool _edgeSelected,
    bool _pointSelected,
    int _pointColorId)
{
    SelectionResult res;

    res.primOrInstanceSel = _primOrInstanceSel;
    res.subprimSel = InitSubprimSelectionResult(
                        _elementSelected, _edgeSelected,
                        _pointSelected, _pointColorId);

    return res;
}


// Decodes the selection buffer encoding scheme for a given mode, and returns
// the selection result.
SelectionResult GetSelectionResult(int mode,
                                   int elementId, int edgeId, int pointId)
{
    SelectionResult res = InitSelectionResult(false, false, false, false, 0);

#if defined(HD_HAS_hdxSelectionBuffer)
    // The hdxSelectionBuffer layout is:
    // [#modes] [per-mode offset] [data mode0] ... [data modeM]
    // [---------header---------]
    // Each mode's data is laid out as:
    // [ prims | points | edges | elements | instance level-N | ... | level 0 ]
    //          <-------- subprims ------->  <----------- instances --------->
    //          <---------------------- per prim ---------------------------->
    // See hdx/SelectionTracker.cpp for details on the encoding scheme.
    int modeOffset = GetIndexOffsetForHighlightMode(mode);
    if (!HighlightModeHasSelection(modeOffset)) {
        // highlight mode has no selected objects (prims/instances/elements)
        return res;
    }

    const int PRIM_SELOFFSETS_HEADER_SIZE = 2;
    const int primId = HdGet_primID();
    int smin = HdGet_hdxSelectionBuffer(modeOffset);
    int smax = HdGet_hdxSelectionBuffer(modeOffset + 1);

    if (primId >= smin && primId < smax) {
        int offset = modeOffset + PRIM_SELOFFSETS_HEADER_SIZE + primId - smin;
        int nextOffset = 0;
        bool sel = false;
        DecodeSelOffset(HdGet_hdxSelectionBuffer(offset),
                        /*out*/sel, /*out*/nextOffset);

        // At this point, sel indicates whether the fragment corresponds to
        // a prim that needs to be fully highlighted, while a non-zero
        // nextOffset indicates whether additional decoding may be done.

        // We don't currently differentiate between prim, instance and
        // subprim selection highlighting (i.e., visually, they look the
        // same), and thus can skip additional decoding if sel is true.
        // We choose not to, for ease of future customization.
        if (nextOffset != 0) {
            // check if instance (or) subprim(s) are selected
            offset = nextOffset;
            sel = sel || IsInstanceSelected(offset, /*out*/nextOffset);

            if (nextOffset != 0) {
                res.subprimSel =
                    GetSubprimSel(nextOffset, elementId, edgeId, pointId);
            }
        }
        res.primOrInstanceSel = sel;
    }
#endif // HD_HAS_hdxSelectionBuffer

    return res;
}

// Returns the logical OR of the inputs
SelectionResult
_CombineWithOr(SelectionResult a, SelectionResult b)
{
    SelectionResult res = InitSelectionResult(false, false, false, false, 0);

    res.primOrInstanceSel = a.primOrInstanceSel || b.primOrInstanceSel;
    res.subprimSel.elementSelected = a.subprimSel.elementSelected ||
                                   b.subprimSel.elementSelected;
    res.subprimSel.edgeSelected = a.subprimSel.edgeSelected ||
                                b.subprimSel.edgeSelected;
    res.subprimSel.pointSelected = a.subprimSel.pointSelected ||
                                 b.subprimSel.pointSelected;
    // pointColorIndex is ignored.
    return res;
}

// Returns the net selection result aggregating the result of each selection
// mode.
SelectionResult GetNetSelectionResult(int elementId, int edgeId, int pointId)
{
    SelectionResult netResult = InitSelectionResult(false, false, false, false, 0);

#if defined(HD_HAS_hdxSelectionBuffer)
    const int numSelectionModes = GetNumSelectionHighlightModes();
    for (int mode = 0; mode < numSelectionModes; mode++) {
        SelectionResult modeResult =
            GetSelectionResult(mode, elementId, edgeId, pointId);
        netResult = _CombineWithOr(modeResult, netResult);
    }
#endif
    return netResult;
}

--- --------------------------------------------------------------------------
-- glsl Selection.Vertex.PointSel
// Mixin for use in the vertex shader stage.
// Decodes the selection buffer to find out if the current vertex (point) is
// selected. This is called from hdSt/shaders/pointId.glslfx
bool IsPointSelected(int pointId)
{
    bool sel = false;
    #if defined(HD_HAS_hdxSelectionBuffer)
    const int numSelectionModes = GetNumSelectionHighlightModes();
    for (int mode = 0; mode < numSelectionModes; mode++) {
        // At the VS stage, we don't carry enough state to determine the
        // elementId and edgeId. So, use fallback values instead.
        SelectionResult res =
            GetSelectionResult(mode, /*elementId*/-1, /*edgeId*/-1, pointId);
        if (res.subprimSel.pointSelected &&
            !HasCustomizedPointColor(res.subprimSel)) {
            sel = true;
            break;
        }
    } // for each highlight mode
    #endif
    return sel;
}

--- --------------------------------------------------------------------------
--  glsl Selection.Geometry.ElementSel
// Mixin for use in the geometry shader stage.
// Helper functions to determine if the element (face) is selected for either
// a given selection mode, or any of the modes.

FORWARD_DECL(int GetElementID()); // code gen

bool IsElementSelected(int mode)
{
    bool sel = false;
    #if defined(HD_HAS_hdxSelectionBuffer)
    const int numSelectionModes = GetNumSelectionHighlightModes();
    if (mode < numSelectionModes) {
        SelectionResult res =
            GetSelectionResult(mode, GetElementID(),
                               /*edgeId*/-1, /*pointId*/-1);
        if (res.subprimSel.elementSelected) {
            sel = true;
        }
    }
    #endif
    return sel;
}

bool IsSelected()
{
    // Edge and point selection aren't relevant at the GS stage.
    SelectionResult res = GetNetSelectionResult(
        GetElementID(), /*edgeId*/-1, /*pointId*/-1);

   return res.primOrInstanceSel || res.subprimSel.elementSelected;
}

--- --------------------------------------------------------------------------
-- glsl Selection.Geometry.WireSelOffset
// Mixin for use in the geometry shader stage for wireframe rendering.
// See comment below.

FORWARD_DECL(int GetMaxNumSelectionHighlightModes());
FORWARD_DECL(bool IsElementSelected(int mode));
FORWARD_DECL(int GetElementID()); // code gen

vec4 ComputeSelectionOffset()
{
    // For wireframe, we only render the edges of faces.  Because we don't have
    // a way to control which face gets rasterized first, if face A is selected
    // and face B is unselected, it's possible for face B to draw over face A
    // and for the highlight to be dropped.

    // To compensate for this, since we're running the geometry shader anyway,
    // add a small polygon offset for selected faces, to push them in front of
    // unselected faces.

    // For doubly-selected faces (e.g. rollover & selection), apply a double
    // offset.  The code below chooses a unique offset for each combination
    // of selection modes by constructing a bitmask of selection modes for
    // this face, and reinterpreting it as an integer.

    int offset = 0;

    int numModes = GetMaxNumSelectionHighlightModes();
    for (int mode = 0; mode < numModes; ++mode) {
        if (IsElementSelected(mode)) {
            offset += (1 << mode);
        }
    }

    // Note: for our base epsilon, we're choosing this somewhat arbitrarily.
    // This is an eye space epsilon, so it's not dependent on distance-from-
    // camera.  Rather, it's dependent on scene scale.  1e-3 works well for
    // our scenes, but if this causes trouble we could use an epsilon based on
    // the eye-space distance of a z-buffer bit flip, which would match the
    // behavior of glPolygonOffset better, at the cost of more math and some
    // extra shader uniforms.
    float eps = 1e-3;

    return vec4(0, 0, offset * eps, 0);
}

--- --------------------------------------------------------------------------
-- glsl Selection.Geometry.WireSelNoOffset

vec4 ComputeSelectionOffset()
{
    return vec4(0);
}

--- --------------------------------------------------------------------------
-- glsl Selection.Fragment
// Mixin for use in the fragment shader stage to determine if the fragment is
// selected.

FORWARD_DECL(int GetElementID()); // generated via codeGen
FORWARD_DECL(int GetPrimitiveEdgeId()); // defined in edgeId.glslfx, or generated via codeGen
FORWARD_DECL(int GetPointId()); // defined in pointId.glslfx

bool IsSelected()
{
    SelectionResult res = GetNetSelectionResult(
        GetElementID(), GetPrimitiveEdgeId(), GetPointId());

   return res.primOrInstanceSel || res.subprimSel.elementSelected;
}

--- --------------------------------------------------------------------------
-- glsl Selection.ComputeColor

#if defined(HD_HAS_hdxSelectionBuffer)
vec4 GetSelectionColor(int mode)
{
    vec4 s = vec4(0,0,0,0);
    // XXX: Make selection colors an array so we can avoid the branching.
    if (mode == 0)
        s = HdGet_selColor();
    else if (mode == 1)
        s = HdGet_selLocateColor();

    return s;
}

// fwd decl fn defined in edgeId.glslfx or generated via codeGen
FORWARD_DECL(float GetSelectedEdgeOpacity());

vec4 GetSubprimSelectionColor(int mode, SubprimSelectionResult res)
{
    vec4 s = GetSelectionColor(mode);
    if (res.edgeSelected) {
        s.a = GetSelectedEdgeOpacity();
    }
    if (res.pointSelected && HasCustomizedPointColor(res)) {
        #if defined(HD_HAS_selectionPointColors)
        s = HdGet_selectionPointColors(res.pointColorId);
        #endif
    }

    return s;
}
#endif

// Fwd declare necessary methods to determine the subprim id of a fragment.
FORWARD_DECL(int GetElementID()); // generated via codeGen
FORWARD_DECL(int GetPrimitiveEdgeId()); // defined in edgeId.glslfx, or generated via codeGen
FORWARD_DECL(int GetPointId()); // defined in pointId.glslfx

// Decodes the selection buffer to find out if the current fragment is from
// a prim/instance/subprim that is selected, applies selection highlighting to
// the incoming color, and returns the resulting color.
vec4 ApplySelectionColor(vec4 color)
{
#if defined(HD_HAS_hdxSelectionBuffer)
    int elementId = GetElementID();
    int edgeId    = GetPrimitiveEdgeId();
    int pointId   = GetPointId();

    const int numSelectionModes = GetNumSelectionHighlightModes();

    bool isSelected = false;
    for (int mode = 0; mode < numSelectionModes; mode++) {
        SelectionResult res =
            GetSelectionResult(mode, elementId, edgeId, pointId);

        if (res.primOrInstanceSel) {
            isSelected = true;
            vec4 s = GetSelectionColor(mode);
            color.rgb = mix(color.rgb, s.rgb, s.a);
        }
        if (IsAnySubprimSelected(res.subprimSel)) {
            isSelected = true;
            vec4 ss = GetSubprimSelectionColor(mode, res.subprimSel);
            color.rgb = mix(color.rgb, ss.rgb, ss.a);
        }
    } // for each highlight mode
    
#endif
    return color;
}
